///////////////////////////////////////////////////////////////////////////
// 1. Includes, macros, etc.
/////////////////////////////////////////////////////////////////////////////

// includes
#include <unistd.h>
#include <ctype.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <tgmath.h>
#include <time.h>
#include <gsl/gsl_cdf.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_errno.h>
#include <gsl/gsl_types.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_interp.h>
#include <gsl/gsl_matrix.h>
#include <omp.h>

// macros: discretization
#define NI 15 // number of sectors
#define NG 1469 // number of goods (previously 1742)
#define NZ 201 // productivity shock grid size
#define NT 100 // simulation length

// macros: paralellization
#ifdef _OPENMP
#define PAR 1
#else
#define PAR 0
#endif

//  macros: output path
#define OUTPUT "output/"
#define INPUT "input/"

// hardcoded params
// time starts in 1971
const int t_reform = 9; // = 1980 - 1971
const int t_wto = 30; // = 2001 - 1971
const int t_1990 = 19; // 1990 - 1971
const int t_data_max = 37; // = 2008 - 1971
const int root_max_iter = 100;
const double root_tol_rel = 1.0e-6;
const double root_tol_abs = 1.0e-6;
const double delta_deriv = 1.0e-9;
const double delta_root = 1.0e-9;
const double policy_tol_abs = 1.0e-10;
const int policy_max_iter = 20000;
const double x_grid_ub_mult = 10.0;
const double x_grid_exp = 1.0;
const double tpu_prob_update_speed = 0.025;
const double coeff_err_tol = 1.0e-2;
const int max_cal_iter = 500;
const double cal_tol=5.0e-2;
const double cal_expart_mult=0.1;
const double cal_exit_mult=0.1;
const double cal_inc_prem_mult=0.03;
const double cal_size_dist_mult=0.03;
const double prob_tol = 1.0e-3;

// print verbose output y/n
const int verbose=1;
int sunk_cost=0;
int calvo=0;
int alt_coeffs=0;
int sens_det=0;
int sens_xi=0;
int ags=0;
int perm_tpu=0;
int iceberg=0;
int ps = 0;
int one_sector=0;
int multi_sector=1;
int cal_probs_flag=0;
int cal_exporter_params_flag=0;
int load_exporter_params_flag=1;

// sign
double sign(double x)
{
  if (x > 0.0) return 1.0;
  if (x < 0.0) return -1.0;
  return x;
}

// initialize all elements of an array to the same numeric value
void set_all_v(double * v, int n, double val)
{
  int i;
  for(i=0; i<n; i++)
    {
      v[i]=val;
    }
}

// macro used for applying above to slices of multidimensional statically allocated arrays
#define SET_ALL_V(v,n,val) ( set_all_v( (double *)(v), (n), (val) ) )

// linspace
void linspace(double lo, double hi, int n, double * v)
{
  double d=(hi-lo)/(n-1.0);
  v[0]=lo;
  int i=0;
  for(i=1;i<n;i++)
    {
      v[i] = v[i-1]+d;
    }
}

// stationary distribution of Markov chain
void markov_stationary_dist(int n, double P[n][n], double p[n])
{
  SET_ALL_V(p,n,1.0/n);
  double diff=-HUGE_VAL;
  do
    {
      diff=-HUGE_VAL;
      
      double tmp[n];
      
      for(int i=0; i<n; i++)
	{
	  tmp[i]=0.0;
	  
	  for(int j=0; j<n; j++)
	    {
	      tmp[i] += P[j][i]*p[j];
	    }
	}
      for(int i=0; i<n; i++)
	{
	  if(fabs(tmp[i]-p[i])>diff)
	    {
	      diff=fabs(tmp[i]-p[i]);
	    }
	}
       for(int i=0; i<n; i++)
	 {
	   p[i]=tmp[i];
	 }
    }
  while(diff>1.0e-11);
}

// horizontal line breaks
void linebreak()
{
  printf("\n////////////////////////////////////////////////////////////////////////////\n");
}

void linebreak2()
{ 
  printf("\n----------------------------------------------------------------------------\n");
}

/////////////////////////////////////////////////////////////////////////////
// 2. Declarations of parameters, grids, and inline functions
/////////////////////////////////////////////////////////////////////////////

// parameters
double W = 0.0; // wage (note: represents normalization of export country GDP per capita relative to representative destination)
double Q = 0.0; // discount factor
double delta0[NI] = {0.0}; // survival rate parameter 1
double delta1[NI] = {0.0}; // survival rate parameter 2
double delta[NI][NZ] = {{0.0}}; // survival rate vector
double theta[NI+1] = {0.0}; // EoS between varieties
double theta_hat[NI] = {0.0}; // = (1/theta)*(theta/(theta-1))^(1-theta)
double theta_hat2[NI] = {0.0}; // = theta*(1/theta)*(theta/(theta-1))^(1-theta)
double sig_z[NI] = {0.0}; // stochastic productivity dispersion
double rho_z[NI] = {0.0}; // stochastic productivity persistence
double mu_e[NI] = {0.0}; // new entrant productivity
double kappa0[NI] = {0.0}; // entry cost
double kappa1[NI] = {0.0}; // continuation cost
double xi[NI] = {0.0}; // iceberg cost in high state
double rho0[NI] = {0.0}; // transition probability from low state
double rho1[NI] = {0.0}; // transition probability from high state

// productivity shock grid
double z_grid[NI][NZ] = {{0.0}}; // grid
double z_hat[NI][NZ] = {{0.0}}; // z^{theta-1} grid
double z_ucond_probs[NI][NZ] = {{0.0}}; // ergodic probabilities
double z_ucond_cumprobs[NI][NZ] = {{0.0}}; // cumulative ergodic probabilities
double z_trans_probs[NI][NZ][NZ] = {{{0.0}}}; // transition probabilities
double z_trans_cumprobs[NI][NZ][NZ] = {{{0.0}}}; // cumultative transition probabilities

// industry-level exporter microdata statistics
double export_participation[NI+1] = {0.0}; // export participation rate
double exit_rate[NI+1] = {0.0}; // exit rate
double incumbent_premium[NI+1] = {0.0}; // relative size of entrants to incumbents
double size_dist[NI+1] = {0.0}; // some moment of the size distribution

// good-level tariff data
int good_2_ind[NG] = {0}; // mapping between goods and industries
double tau_applied[NG][NT] = {{0.0}}; // trade cost before liberalization
double tau_nntr[NG] = {0.0}; // trade cost before liberalization
double gap[NG] = {0.0}; // NNTR gap

// aggregate trade policy
double tpu_trans_mat[2][2][NT] = {{{0.0}}}; // probability of switching between NNTR and MFN
double iceberg_val[NG][NT] = {{0.0}}; //  iceberg costs after liberalization
double frac_1980 = 0.5; // fraction of firms that get MFN tariffs in 1980... needed to match SR ECM elasticity

// PS probability
double tpu_prob_ps = 0.25;

// load data
int load_input_data()
{
  // sector-level data
  FILE * file = fopen(INPUT "inputs_calibration_sector.csv","r");
  if(!file)
    {
      printf("Failed to open input file with sector-level data!\n");
      return 1;
    }

  int ii;
  double theta_, partic_, exit_, inc_prem_, size_dist_;
  int cnt=0;  
  while(cnt<=NI && fscanf(file,"%d %lf %lf %lf %lf %lf",&ii ,&theta_, &partic_, &exit_, &inc_prem_, &size_dist_) == 6)
    {
      cnt+=1;
      
      if(ii<0 || ii>NI || gsl_isnan(theta_) || gsl_isnan(partic_) ||
	 gsl_isnan(exit_) || gsl_isinf(inc_prem_) || gsl_isinf(size_dist_))
	{
	  printf("Bad value in sector-level data!\n");
	  fclose(file);
	  return 1;
	}
      
      theta[ii] = theta_;
      export_participation[ii] = partic_;
      exit_rate[ii] = exit_;
      incumbent_premium[ii] = inc_prem_;
      size_dist[ii] = size_dist_;      
    }
  fclose(file);

  // good-level data
  file = fopen(INPUT "inputs_calibration_g.csv","r");
  if(!file)
    {
      printf("Failed to open input file with good-level data!\n");
      return 1;
    }

  int ig, g2i_;
  double nntr_, gap_;
  cnt=0;  
  while(cnt<NG && fscanf(file,"%d %d %lf %lf",&ig ,&g2i_, &nntr_, &gap_) == 4)
    {
      cnt+=1;
      
      if(ig<0 || ig>NG-1 || g2i_<0 || g2i_>NI-1 || gsl_isnan(nntr_) || gsl_isnan(gap_))
	{
	  printf("Bad value in good-level data!\n");
	  fclose(file);
	  return 1;
	}

      good_2_ind[ig] = g2i_;
      tau_nntr[ig] = 1.0+nntr_;
      gap[ig] = gap_;
    }
  fclose(file);
  
  // good-year level data
  file = fopen(INPUT "inputs_calibration_gt.csv","r");
  if(!file)
    {
      printf("Failed to open file with good-year level data!\n");
      return 1;
    }

  int t;
  double  mfn_;
  cnt=0;      
  while(fscanf(file,"%d %d %lf",&ig,&t,&mfn_) == 3)
    {
      cnt+=1;
      
      if(ig<0 || ig>NG-1 || t<1974 || t>2008 || gsl_isnan(mfn_))
	{
	  printf("Bad value in good-year level data! %d %d %0.6f\n",ig,t,mfn_);
	  fclose(file);
	  return 1;
	}
      
      tau_applied[ig][t-1971] = 1.0 + mfn_;
    }
  fclose(file);

  if(cnt != (NG*(t_data_max-3+1)))
    {
      printf("Failed to load data cnt = %d, should be %d!\n",cnt,NG*(t_data_max-3+1));
      return 1;
    }
  
  for(int ig=0; ig<NG; ig++)
    {
      for(int t=0; t<3; t++)
	{
	  tau_applied[ig][t] = tau_applied[ig][3];
	}
    }
      
  for(int t=t_data_max+1; t<NT; t++)
    {
      for(int ig=0; ig<NG; ig++)
	{
	  tau_applied[ig][t] = tau_applied[ig][t-1];
	}
    }
  
  for(int t=0; t<NT; t++)
    {
      for(int ig=0; ig<NG; ig++)
	{
	  double tau_ = tau_applied[ig][t];
	  double nntr_ = tau_nntr[ig];
	  if(gsl_isnan(tau_) || gsl_isnan(nntr_) || gsl_isnan(gap_) ||
	     gsl_isinf(tau_) || gsl_isinf(nntr_) || gsl_isinf(gap_) ||
	     tau_<0.999 || nntr_ < 0.99)
		
	    {
	      printf("Bad tariff data!\n");
	      printf("%d %d %0.16f %0.16f\n",t,ig,tau_,nntr_);
	      return 1;
	    }
	}
    }

  // initial guess for probabilities
  // 0: NNTR
  // 1: MFN
  if(ps==0)
    {
      for(int t=0; t<NT; t++)
	{
	  double frac = (double)(t-t_reform)/ ((double)(t_data_max-t_reform));
      
	  tpu_trans_mat[0][0][t] = 0.95;
	  tpu_trans_mat[0][1][t] = 0.05;

	  if(t<t_reform)
	    tpu_trans_mat[1][0][t] = 0.8;
	  else if(t<=t_data_max)
	    tpu_trans_mat[1][0][t] = 0.8 * (1.0-sqrt(frac)) + sqrt(frac) * 0.01;
	  else
	    tpu_trans_mat[1][0][t] = tpu_trans_mat[1][0][t-1];
      
	  tpu_trans_mat[1][1][t] = 1.0-tpu_trans_mat[1][0][t];
	}

      if(one_sector==1)
	file = fopen(OUTPUT "tpuprobs_one_sector.txt","r");
      else if(one_sector==2)
	file = fopen(OUTPUT "tpuprobs_avg_sector.txt","r");
      else if(calvo==1)
	file = fopen(OUTPUT "tpuprobs_calvo.txt","r");
      else if(sens_xi==1)
	file = fopen(OUTPUT "tpuprobs_lo_xi.txt","r");
      else if(sens_xi==2)
	file = fopen(OUTPUT "tpuprobs_hi_xi.txt","r");
      else if(sens_xi==3)
	file = fopen(OUTPUT "tpuprobs_lo_rhoxi.txt","r");
      else if(sens_xi==4)
	file = fopen(OUTPUT "tpuprobs_hi_rhoxi.txt","r");
      else if(alt_coeffs==0 && ps==0)
	{
	  if(perm_tpu==1)
	    file = fopen(OUTPUT "tpuprobs_permtpu.txt","r");
	  else
	    file = fopen(OUTPUT "tpuprobs_baseline.txt","r");
	}
      else if(alt_coeffs==5)
	file = fopen(OUTPUT "tpuprobs_ci_lower.txt","r");
      else if(alt_coeffs==6)
	file = fopen(OUTPUT "tpuprobs_ci_upper.txt","r");
  
      if(!file)
	{
	  printf("Failed to open file with probabilities!\n");
	  return 1;
	}
      else
	{
	  double prob;
	  int cnt=0;
	  for(int t=2; t<t_reform; t++)
	    {
	      cnt += fscanf(file,"%lf",&prob);
	      tpu_trans_mat[0][1][t] = prob;
	    }
	  for(int t=t_reform; t<t_data_max; t++)
	    {
	      cnt += fscanf(file,"%lf",&prob);
	      tpu_trans_mat[1][0][t] = prob;
	    }
	  fclose(file);

	  if(cnt != t_data_max-3+1)
	    {
	      printf("Failed to load probabilities!\n");
	      return 1;
	    }

	  //tpu_trans_mat[0][1][3] = 0.15;
	  for(int t=0; t<3; t++)
	    {
	      tpu_trans_mat[0][1][t] = tpu_trans_mat[0][1][3];
	    }
	  for(int t=t_reform; t<NT; t++)
	    {
	      tpu_trans_mat[0][1][t] = tpu_trans_mat[0][1][3];
	    }
      
	  for(int t=0; t<t_reform; t++)
	    {
	      tpu_trans_mat[1][0][t] = tpu_trans_mat[1][0][t_reform];
	    }
	  for(int t=t_data_max; t<NT; t++)
	    {
	      tpu_trans_mat[1][0][t] = tpu_trans_mat[1][0][t_data_max-1];
	    }

	  for(int t=0; t<NT; t++)
	    {
	      tpu_trans_mat[0][0][t] = 1.0-tpu_trans_mat[0][1][t];
	      tpu_trans_mat[1][1][t] = 1.0-tpu_trans_mat[1][0][t];
	    }
	}
    }

  return 0;
}

// assigned parameters and initial guesses
int init_params()
{
  // params constant to all industries
  W = 1.0;
  Q = 0.96;

  if(multi_sector==1)
    {
      frac_1980 = 0.69;
    }
  if(one_sector==1)
    {
      frac_1980 = 0.67;
    }
  else if(one_sector==2)
    {
      frac_1980 = 0.7;
    }

  // model-invariant params
  for(int ii=0; ii<NI; ii++)
    {
      // if we are assuming same parameters across all sectors, set demand elasticity
      // to the last sector's, which is the cross-sector average
      if(one_sector==1)
	theta[ii] = theta[NI];
      else if(one_sector==2)
	theta[ii] = theta[NI-1];
      
      theta_hat[ii] = (1.0/theta[ii]) * pow(theta[ii]/(theta[ii]-1.0),1.0-theta[ii]);
      theta_hat2[ii] = theta[ii] * theta_hat[ii]; 
      rho_z[ii] =  0.65;
      mu_e[ii] = 0.0;
      delta0[ii] = 21.04284098;

      if(multi_sector==1)
	{
	  rho0[ii]=0.9075;
	  rho1[ii]=rho0[ii];
	}  
      else if(one_sector==1)
	{
	  rho0[ii]=0.89;
	  rho1[ii]=rho0[ii];
	}
      else if(one_sector==2)
	{
	  rho0[ii]=0.92;
	  rho1[ii]=rho0[ii];
	}

      if(calvo==0)
	delta1[ii] = 0.02258301;
      else
	{
	  delta1[ii] = 0.16;
	  rho_z[ii] = 0.65;
	}
    }      
      
  // initial guesses for exporter parames if we are recalibrating
  if(load_exporter_params_flag==0)
    {
      if(sunk_cost==0 && calvo==0)
	{
	  for(int ii=0; ii<NI; ii++)
	    {
	      sig_z[ii] = 1.32;
	      kappa0[ii]=0.5;
	      kappa1[ii]=0.3;
	      xi[ii]=2.6;	     
	    }
	}
      // sunk cost model
      else if(sunk_cost==1)
	{
	  for(int ii=0; ii<NI; ii++)
	    {
	      sig_z[ii] = 1.32;
	      kappa0[ii] = 2.2;
	      kappa1[ii] = 0.55;
	      xi[ii]=1.0;
	    }
	}
      // calco model
      else if(calvo==1)
	{
	  for(int ii=0; ii<NI; ii++)
	    {
	      sig_z[ii] = 1.32;
	      kappa0[ii] = 3.25;
	      kappa1[ii] = 0.0;
	      xi[ii]=1.0;
	    }
	}
    }
  
  for(int ig=0; ig<NG; ig++)
    {
      for(int t=0; t<NT; t++)
	{
	  iceberg_val[ig][t] = 1.0;
	}
    }

  return 0;
}

// discretization of productivity shock process
void discretize_z(int ii)
{
  int n = NZ;
  double inprob = 1.0e-8;
  double lo = gsl_cdf_ugaussian_Pinv(inprob)*sig_z[ii]*1.5;
  double hi = -gsl_cdf_ugaussian_Pinv(inprob)*sig_z[ii]*1.5;
  double ucond_std = sqrt(sig_z[ii]*sig_z[ii]/(1.0-rho_z[ii]*rho_z[ii]));
  double d = (hi-lo)/(n-1.0);
  linspace(lo,hi,n,z_grid[ii]);
  
  for(int iz=0; iz<n; iz++)
    {
      double x = z_grid[ii][iz];

      double sum=0.0;
      for(int izp=0; izp<n; izp++)
	{
	  double y = z_grid[ii][izp];	  
	  z_trans_probs[ii][iz][izp] = (gsl_cdf_ugaussian_P( (y + d/2.0 - rho_z[ii]*x) / sig_z[ii] ) -
					gsl_cdf_ugaussian_P( (y - d/2.0 - rho_z[ii]*x) / sig_z[ii] ));
	  
	  sum += z_trans_probs[ii][iz][izp];
	}
      for(int izp=0; izp<n; izp++)
	{
	  z_trans_probs[ii][iz][izp] = z_trans_probs[ii][iz][izp]/sum;
	}
    }

  double sum=0.0;
  for(int iz=0; iz<n; iz++)
    {
      double x = z_grid[ii][iz];
      
      z_ucond_probs[ii][iz] = (gsl_cdf_ugaussian_P( (x + mu_e[ii] + d/2.0) / ucond_std ) -
			       gsl_cdf_ugaussian_P( (x + mu_e[ii] - d/2.0) / ucond_std ));
      
      sum += z_ucond_probs[ii][iz];
    }
  for(int iz=0; iz<n; iz++)
    {
      z_ucond_probs[ii][iz] = z_ucond_probs[ii][iz]/sum;
    }

  sum=0.0;
  for(int iz=0; iz<n; iz++)
    {
      z_grid[ii][iz] = exp(z_grid[ii][iz]);
      z_hat[ii][iz] = z_grid[ii][iz];
      sum += z_ucond_probs[ii][iz];
      z_ucond_cumprobs[ii][iz] = sum;

      double sum2=0.0;
      for(int izp=0; izp<n; izp++)
	{
	  sum2 += z_trans_probs[ii][iz][izp];
	  z_trans_cumprobs[ii][iz][izp] = sum2;
	}
    }
}

// survival probability vector
void calc_death_probs(int ii)
{
  for(int iz=0; iz<NZ; iz++)
    {
      double death_prob=fmax(0.0,fmin(exp(-delta0[ii]*z_hat[ii][iz])+delta1[ii],1.0));
      delta[ii][iz] = 1.0-death_prob;
    }
}

/////////////////////////////////////////////////////////////////////////////
// 3. Dynamic program: deterministic
/////////////////////////////////////////////////////////////////////////////

double V_nntr[NG][NZ][3] = {{{0.0}}}; // NNTR
double V_ref_det[NG][NZ][3][NT] = {{{{0.0}}}}; // perfect foresight
double EV_det[NG][NZ][3] = {{{0.0}}}; // continuation value storage

int gex_nntr[NG][NZ][3] = {{{0}}}; // NNTR
int gex_ref_det[NG][NZ][3][NT] = {{{{0}}}}; // perfect foresight

int policy_solved_flag[NG] = {0};

// initial guess for value functions
void init_dp_objs_det()
{
  for(int ig=0; ig<NG; ig++)
    {
      int ii = good_2_ind[ig];
      
      for(int iz=0; iz<NZ; iz++)
	{ 
	  double pi_hat = theta_hat[ii] * pow(tau_nntr[ig],-theta[ii]);
	  V_nntr[ig][iz][0] = 0.0;
	  V_nntr[ig][iz][1] = pi_hat*pow(xi[ii],1.0-theta[ii])*z_hat[ii][iz]/Q;
	  V_nntr[ig][iz][2] = pi_hat*z_hat[ii][iz]/Q;
	  
	  for(int t=0; t<NT; t++)
	    {
	      pi_hat = theta_hat[ii] * pow(tau_applied[ig][t],-theta[ii]) * pow(iceberg_val[ig][t],1.0-theta[ii]);
	      V_ref_det[ig][iz][0][t] = 0.0;
	      V_ref_det[ig][iz][1][t] = pi_hat*pow(xi[ii],1.0-theta[ii])*z_hat[ii][iz]/Q;
	      V_ref_det[ig][iz][2][t] = pi_hat*z_hat[ii][iz]/Q;
	    }
	}
    }
}

// steady state Bellman
void calc_EV_det(int ig, int t, int flag)
{
  int ii = good_2_ind[ig];
  
  for(int z=0; z<NZ; z++)
    {
      for(int e=0; e<3; e++)
	{
	  if(flag==0)
	    {
	      EV_det[ig][z][e]=0.0;
	    }
	  else
	    {
	      EV_det[ig][z][e]=0.0;
	    }
	  
	  for(int zp=0; zp<NZ; zp++)
	    {
	      if(z_trans_probs[ii][z][zp]>1.0e-11)
		{
		  if(flag==0)
		    {
		      EV_det[ig][z][e] += Q*delta[ii][z]*
			V_nntr[ig][zp][e]*z_trans_probs[ii][z][zp];
		    }
		  else
		    {
		      EV_det[ig][z][e] += Q*delta[ii][z]*
			V_ref_det[ig][zp][e][t]*z_trans_probs[ii][z][zp];
		    }
		}
	    }
	} 
    }

}

void iterate_policies_det(int ig, int t, int flag,
			  double * maxdiff, int imaxdiff[3])
{
  int ii = good_2_ind[ig];
  
  *maxdiff=-HUGE_VAL;

  if(t<NT-1)
    calc_EV_det(ig,t+1,flag);
  else
    calc_EV_det(ig,t,flag);
 
  for(int z=0; z<NZ; z++)
    {
      int * gex1, * gex2, * gex3;
      if(flag==0)
	{
	  gex1 = &(gex_nntr[ig][z][0]);
	  gex2 = &(gex_nntr[ig][z][1]);
	  gex3 = &(gex_nntr[ig][z][2]);
	}
      else
	{
	  gex1 = &(gex_ref_det[ig][z][0][t]);
	  gex2 = &(gex_ref_det[ig][z][1][t]);
	  gex3 = &(gex_ref_det[ig][z][2][t]);
	}
    

      // applied
      if(EV_det[ig][z][0] < EV_det[ig][z][1] - kappa0[ii])
	{
	  *gex1 = 1;
	}
      else
	{
	  *gex1 = 0;
	}

      if(EV_det[ig][z][0] <
	 rho0[ii]*EV_det[ig][z][1] + (1.0-rho0[ii])*EV_det[ig][z][2] - kappa1[ii])
	{
	  *gex2 = 1;
	}
      else
	{
	  *gex2 = 0;
	}
	  
      if(EV_det[ig][z][0] <
	 (1.0-rho1[ii])*EV_det[ig][z][1] + rho1[ii]*EV_det[ig][z][2] - kappa1[ii])
	{
	  *gex3 = 1;
	}
      else
	{
	  *gex3 = 0;
	}
      
      // update continuation values and check convergence ---------------
      double pi=0.0;
      if(flag==0)
	pi = theta_hat[ii] * pow(tau_nntr[ig],-theta[ii]) * z_hat[ii][z];
      else if(t<t_reform)
	pi = theta_hat[ii] * pow(tau_nntr[ig],-theta[ii]) * pow(iceberg_val[ig][t],1.0-theta[ii]) * z_hat[ii][z];
      else
	pi = theta_hat[ii] * pow(tau_applied[ig][t],-theta[ii]) * pow(iceberg_val[ig][t],1.0-theta[ii]) * z_hat[ii][z];
      
      double tmp0 = fmax(EV_det[ig][z][0], EV_det[ig][z][1] - kappa0[ii]);
      
      double tmp1 = pi*pow(xi[ii],1.0-theta[ii]) +
	fmax(EV_det[ig][z][0],
	     rho0[ii]*EV_det[ig][z][1] +
	     (1.0-rho0[ii])*EV_det[ig][z][2] - kappa1[ii]);

      double tmp2 = pi +
	fmax(EV_det[ig][z][0],
	     (1.0-rho1[ii])*EV_det[ig][z][1] +
	     rho1[ii]*EV_det[ig][z][2] - kappa1[ii]);

	
      double diff0=0, diff1=0, diff2=0;
      if(flag==0)
	{
	  diff0 = fabs(tmp0-V_nntr[ig][z][0]);
	  diff1 = fabs(tmp1-V_nntr[ig][z][1]);
	  diff2 = fabs(tmp2-V_nntr[ig][z][2]);
	}
      else
	{
	  diff0 = fabs(tmp0-V_ref_det[ig][z][0][t]);
	  diff1 = fabs(tmp1-V_ref_det[ig][z][1][t]);
	  diff2 = fabs(tmp2-V_ref_det[ig][z][2][t]);
	}

      if(diff0>*maxdiff)
	{
	  *maxdiff=diff0;
	  imaxdiff[0]=z;
	  imaxdiff[1]=0;
	  imaxdiff[2]=1;
	}

      if(diff1>*maxdiff)
	{
	  *maxdiff=diff1;
	  imaxdiff[0]=z;
	  imaxdiff[1]=1;
	  imaxdiff[2]=1;
	}
	  
      if(diff2>*maxdiff)
	{
	  *maxdiff=diff2;
	  imaxdiff[0]=z;
	  imaxdiff[1]=2;
	  imaxdiff[2]=1;
	}

      if(flag==0)
	{
	  V_nntr[ig][z][0] = tmp0;
	  V_nntr[ig][z][1] = tmp1;
	  V_nntr[ig][z][2] = tmp2;
	}
      else
	{
	  V_ref_det[ig][z][0][t] = tmp0;
	  V_ref_det[ig][z][1][t] = tmp1;
	  V_ref_det[ig][z][2][t] = tmp2;
	}
    }
}

// solve policy function for industry i
int solve_policies_det(int ig)
{
  double maxdiff = 999;
  int imaxdiff[3];

  // first do SS policy for NNTR rates
  int iter=0;
  do
    {
      iter++;
      iterate_policies_det(ig,NT-1,0,&maxdiff,imaxdiff);
    }
  while(maxdiff>policy_tol_abs && iter < policy_max_iter);

  if(iter==policy_max_iter)
    {
      printf("\tValue function iteration failed for industry %d! Diff = %0.4g\n",ig,maxdiff);
      return 1;
    }

  // now do SS policy for very last period of MFN rates
  do
    {
      iter++;
      iterate_policies_det(ig,NT-1,1,&maxdiff,imaxdiff);
    }
  while(maxdiff>policy_tol_abs && iter < policy_max_iter);

  if(iter==policy_max_iter)
    {
      printf("\tValue function iteration failed for industry %d! Diff = %0.4g\n",ig,maxdiff);
      return 1;
    }

  // now iterate backwards to get path of MFN policies
  for(int t=NT-2; t>=0; t--)
    {
      iterate_policies_det(ig,t,1,&maxdiff,imaxdiff);
    }

  return 0;
}

// solve policies for all industries in parallel
int solve_policies2_det()
{
  if(verbose)
    printf("\nSolving deterministic dynamic program...\n");

  time_t start, stop;
  time(&start);

  int cnt=0;
  
#ifdef _OPENMP
#pragma omp parallel for
#endif
  for(int ig=0; ig<NG; ig++)
    {
      policy_solved_flag[ig] = solve_policies_det(ig);
    }

  for(int ig=0; ig<NG; ig++)
    cnt += 1-policy_solved_flag[ig];

  time(&stop);
  
  if(verbose)
    {
      printf("Finished %0.0f seconds. %d/%d converged.\n",difftime(stop,start),cnt,NG);
    }
  
  return 0;  
}

/////////////////////////////////////////////////////////////////////////////
// 4a. Dynamic program: Markov process
/////////////////////////////////////////////////////////////////////////////

double V_markov[NG][NZ][3][2][NT] = {{{{{0.0}}}}}; // Markov process
double EV_markov[NG][NZ][3][2] = {{{{0.0}}}}; // Markov process
int gex_markov[NG][NZ][3][2][NT] = {{{{{0}}}}}; // Markov process

// initial guess for value functions
void init_dp_objs_markov()
{
  for(int ig=0; ig<NG; ig++)
    {
      int ii = good_2_ind[ig];
      
      for(int iz=0; iz<NZ; iz++)
	{
	  double pi_hat = theta_hat[ii] * pow(tau_applied[ig][NT-1],-theta[ii]);
	  V_markov[ig][iz][0][1][NT-1] = 0.0;
	  V_markov[ig][iz][1][1][NT-1] = pi_hat*pow(xi[ii],1.0-theta[ii])*z_hat[ii][iz]/Q;
	  V_markov[ig][iz][2][1][NT-1] = pi_hat*z_hat[ii][iz]/Q;
	  
	  pi_hat = theta_hat[ii] * pow(tau_nntr[ig],-theta[ii]);
	  V_markov[ig][iz][0][0][NT-1] = 0.0;
	  V_markov[ig][iz][1][0][NT-1] = pi_hat*pow(xi[ii],1.0-theta[ii])*z_hat[ii][iz]/Q;
	  V_markov[ig][iz][2][0][NT-1] = pi_hat*z_hat[ii][iz]/Q;
	}
    }
}

// steady state Bellman
void calc_EV_markov(int ig, int t, int tp)
{
  int ii = good_2_ind[ig];
  
  for(int z=0; z<NZ; z++)
    {
      for(int e=0; e<3; e++)
	{
	  for(int p=0; p<2; p++)
	    {
	      EV_markov[ig][z][e][p]=0.0;
	      for(int zp=0; zp<NZ; zp++)
		{
		  if(z_trans_probs[ii][z][zp]>1.0e-11)
		    {
		      for(int pp=0; pp<2; pp++)
			{
			  EV_markov[ig][z][e][p] += Q*delta[ii][z]*
			    V_markov[ig][zp][e][pp][t]*
			    z_trans_probs[ii][z][zp]*tpu_trans_mat[p][pp][tp];
			}
		    }
		}
	    }
	} 
    }

}

void iterate_policies_markov(int ig, int t, int tp, double * maxdiff, int imaxdiff[3])
{
  int ii = good_2_ind[ig];
  
  *maxdiff=-HUGE_VAL;

  if(t==NT-1)
    calc_EV_markov(ig,t,tp);
  else
    calc_EV_markov(ig,t+1,tp);
  
  for(int z=0; z<NZ; z++)
    {
      // first compute policy functions -------------------
      
      // NNTR = 0, MFN = 1
      for(int p=0; p<2; p++)
	{
	  if(EV_markov[ig][z][0][p] < EV_markov[ig][z][1][p] - kappa0[ii])
	    {
	      gex_markov[ig][z][0][p][t] = 1;
	    }
	  else
	    {
	      gex_markov[ig][z][0][p][t] = 0;
	    }

	  if(EV_markov[ig][z][0][p] <
	     rho0[ii]*EV_markov[ig][z][1][p] +
	     (1.0-rho0[ii])*EV_markov[ig][z][2][p] - kappa1[ii])
	    {
	      gex_markov[ig][z][1][p][t] = 1;
	    }
	  else
	    {
	      gex_markov[ig][z][1][p][t] = 0;
	    }
	  
	  if(EV_markov[ig][z][0][p] <
	     (1.0-rho1[ii])*EV_markov[ig][z][1][p] +
	     rho1[ii]*EV_markov[ig][z][2][p] - kappa1[ii])
	    {
	      gex_markov[ig][z][2][p][t] = 1;
	    }
	  else
	    {
	      gex_markov[ig][z][2][p][t] = 0;
	    }
	}
      
      // update continuation values and check convergence ---------------

      // autarky
      for(int p=0; p<2; p++)
	{
	  double tmp0=0.0;
	  double tmp1=0.0;
	  double tmp2=0.0;
	  
	  double pi = 0.0;
	  if(p==0)
	    pi = theta_hat[ii] * pow(tau_nntr[ig],-theta[ii]) * z_hat[ii][z];
	  else if(t<t_reform)
	    pi = theta_hat[ii] * pow(tau_applied[ig][t_reform],-theta[ii]) * z_hat[ii][z];
	  else
	    pi = theta_hat[ii] * pow(tau_applied[ig][t],-theta[ii]) * z_hat[ii][z];
		

	  tmp0 = fmax(EV_markov[ig][z][0][p],
		      EV_markov[ig][z][1][p] - kappa0[ii]);
      
	  tmp1 = pi*pow(xi[ii],1.0-theta[ii]) +
	    fmax(EV_markov[ig][z][0][p],
		 rho0[ii]*EV_markov[ig][z][1][p] +
		 (1.0-rho0[ii])*EV_markov[ig][z][2][p] - kappa1[ii]);

	  tmp2 = pi +
	    fmax(EV_markov[ig][z][0][p],
		 (1.0-rho1[ii])*EV_markov[ig][z][1][p] +
		 rho1[ii]*EV_markov[ig][z][2][p] - kappa1[ii]);
	  
	  double diff0 = fabs(tmp0-V_markov[ig][z][0][p][t]);
	  double diff1 = fabs(tmp1-V_markov[ig][z][1][p][t]);
	  double diff2 = fabs(tmp2-V_markov[ig][z][2][p][t]);

	  if(diff0>*maxdiff)
	    {
	      *maxdiff=diff0;
	      imaxdiff[0]=z;
	      imaxdiff[1]=0;
	      imaxdiff[2]=0;
	    }
	  
	  if(diff1>*maxdiff)
	    {
	      *maxdiff=diff1;
	      imaxdiff[0]=z;
	      imaxdiff[1]=1;
	      imaxdiff[2]=0;
	    }
	  
	  if(diff2>*maxdiff)
	    {
	      *maxdiff=diff2;
	      imaxdiff[0]=z;
	      imaxdiff[1]=2;
	      imaxdiff[2]=0;
	    }

	  V_markov[ig][z][0][p][t] = tmp0;
	  V_markov[ig][z][1][p][t] = tmp1;
	  V_markov[ig][z][2][p][t] = tmp2;
	}
    }
}

// solve policy function for industry i
int solve_policies_markov(int ig)
{
  double maxdiff = 999;
  int imaxdiff[3];

  // in the ratex case, where firms know the path of the transition matrix in advance,
  // the approach is standard
  if(perm_tpu==0)
    {
      // first do SS policies for pre-1980 and NNTR rates
      int iter=0;
      do
	{
	  iter++;
	  iterate_policies_markov(ig,NT-1,NT-1,&maxdiff,imaxdiff);
	}
      while(maxdiff>policy_tol_abs && iter < policy_max_iter);

      if(iter==policy_max_iter)
	{
	  printf("\tValue function iteration failed for industry %d! Diff = %0.4g\n",ig,maxdiff);
	  return 1;
	}

      // now iterate backwards
      for(int t=NT-2; t>=0; t--)
	{
	  iterate_policies_markov(ig,t,t,&maxdiff,imaxdiff);
	}
    }
  // in the case where in each period t, firms think that period's transition matrix will
  // last forever, things are a lot more complicated
  else
    {
      // for each period t_outer (which is associated with a specific transition matrix)...
      for(int t_outer=0; t_outer<NT; t_outer++)
	{
	  // first solve terminal-period SS policies with the 
	  int iter=0;
	  do
	    {
	      iter++;
	      iterate_policies_markov(ig,NT-1,t_outer,&maxdiff,imaxdiff);
	    }
	  while(maxdiff>policy_tol_abs && iter < policy_max_iter);
	  
	  if(iter==policy_max_iter)
	    {
	      printf("\tValue function iteration failed for industry %d! Diff = %0.4g\n",ig,maxdiff);
	      return 1;
	    }

	  // and then iterate backward to period t_outer
	  for(int t_inner=NT-2; t_inner>=t_outer; t_inner--)
	    {
	      iterate_policies_markov(ig,t_inner,t_outer,&maxdiff,imaxdiff);
	    }
	}
    }

  return 0;
}

// solve policies for all industries in parallel
int solve_policies2_markov()
{
  if(verbose)
    printf("\nSolving stochastic dynamic program with time-varying transition probabilities...\n");

  time_t start, stop;
  time(&start);

  int cnt=0;
  
#ifdef _OPENMP
#pragma omp parallel for
#endif
  for(int ig=0; ig<NG; ig++)
    {
      policy_solved_flag[ig] = solve_policies_markov(ig);
      cnt += policy_solved_flag[ig];
    }

  time(&stop);
  
  if(verbose)
    {
      printf("Finished dynamic programs in %0.0f seconds. %d failed to converge.\n",difftime(stop,start),cnt);
    }
  
  return 0;  
}

/////////////////////////////////////////////////////////////////////////////
// 4b. Dynamic program: PS/HL with constant prob. of losing MFN
/////////////////////////////////////////////////////////////////////////////

double V_PS[NG][NZ][3][NT] = {{{{0.0}}}}; // Markov process
double EV_PS[NG][NZ][3] = {{{0.0}}}; // Markov process
int gex_PS[NG][NZ][3][NT] = {{{{0}}}}; // Markov process

// initial guess for value functions
void init_dp_objs_PS()
{
  for(int ig=0; ig<NG; ig++)
    {
      int ii = good_2_ind[ig];
      
      for(int iz=0; iz<NZ; iz++)
	{
	  double pi_hat = theta_hat[ii] * pow(tau_applied[ig][NT-1],-theta[ii]);
	  V_PS[ig][iz][0][NT-1] = 0.0;
	  V_PS[ig][iz][1][NT-1] = pi_hat*pow(xi[ii],1.0-theta[ii])*z_hat[ii][iz]/Q;
	  V_PS[ig][iz][2][NT-1] = pi_hat*z_hat[ii][iz]/Q;
	}
    }
}

// steady state Bellman
void calc_EV_PS(int ig, int t, double prob)
{
  int ii = good_2_ind[ig];
  
  for(int z=0; z<NZ; z++)
    {
      for(int e=0; e<3; e++)
	{
	  EV_PS[ig][z][e]=0.0;
	  for(int zp=0; zp<NZ; zp++)
	    {
	      if(z_trans_probs[ii][z][zp]>1.0e-11)
		{
		  EV_PS[ig][z][e] += Q*delta[ii][z]*z_trans_probs[ii][z][zp]*
		    (V_PS[ig][zp][e][t]*(1.0-prob) + V_nntr[ig][zp][e]*prob);
		}
	    }
	}
    }

}

void iterate_policies_PS(int ig, int t, double prob, double * maxdiff, int imaxdiff[3])
{
  int ii = good_2_ind[ig];
  
  *maxdiff=-HUGE_VAL;

  if(t==NT-1)
    calc_EV_PS(ig,t,prob);
  else
    {
      calc_EV_PS(ig,t+1,prob);
    }
  
  for(int z=0; z<NZ; z++)
    {
      // first compute policy functions -------------------
      
      if(EV_PS[ig][z][0] < EV_PS[ig][z][1] - kappa0[ii])
	{
	  gex_PS[ig][z][0][t] = 1;
	}
      else
	{
	  gex_PS[ig][z][0][t] = 0;
	}
      
      if(EV_PS[ig][z][0] <
	 rho0[ii]*EV_PS[ig][z][1] +
	 (1.0-rho0[ii])*EV_PS[ig][z][2] - kappa1[ii])
	{
	  gex_PS[ig][z][1][t] = 1;
	}
      else
	{
	  gex_PS[ig][z][1][t] = 0;
	}
      
      if(EV_PS[ig][z][0] <
	 (1.0-rho1[ii])*EV_PS[ig][z][1] +
	 rho1[ii]*EV_PS[ig][z][2] - kappa1[ii])
	{
	  gex_PS[ig][z][2][t] = 1;
	}
      else
	{
	  gex_PS[ig][z][2][t] = 0;
	}
      
      // update continuation values and check convergence ---------------
      
      // autarky
      double tmp0=0.0;
      double tmp1=0.0;
      double tmp2=0.0;
      
      double pi = 0.0;
      if(t<t_reform)
	pi = theta_hat[ii] * pow(tau_nntr[ig],-theta[ii]) * z_hat[ii][z];
      else
	pi = theta_hat[ii] * pow(tau_applied[ig][t],-theta[ii]) * z_hat[ii][z];
      
      
      tmp0 = fmax(EV_PS[ig][z][0],
		  EV_PS[ig][z][1] - kappa0[ii]);
      
      tmp1 = pi*pow(xi[ii],1.0-theta[ii]) +
	fmax(EV_PS[ig][z][0],
	     rho0[ii]*EV_PS[ig][z][1] +
	     (1.0-rho0[ii])*EV_PS[ig][z][2] - kappa1[ii]);
      
      tmp2 = pi +
	fmax(EV_PS[ig][z][0],
	     (1.0-rho1[ii])*EV_PS[ig][z][1] +
	     rho1[ii]*EV_PS[ig][z][2] - kappa1[ii]);
	  
      double diff0 = fabs(tmp0-V_PS[ig][z][0][t]);
      double diff1 = fabs(tmp1-V_PS[ig][z][1][t]);
      double diff2 = fabs(tmp2-V_PS[ig][z][2][t]);

      if(diff0>*maxdiff)
	{
	  *maxdiff=diff0;
	  imaxdiff[0]=z;
	  imaxdiff[1]=0;
	  imaxdiff[2]=0;
	}
      
      if(diff1>*maxdiff)
	{
	  *maxdiff=diff1;
	  imaxdiff[0]=z;
	  imaxdiff[1]=1;
	  imaxdiff[2]=0;
	}
      
      if(diff2>*maxdiff)
	{
	  *maxdiff=diff2;
	  imaxdiff[0]=z;
	  imaxdiff[1]=2;
	  imaxdiff[2]=0;
	}
      
      V_PS[ig][z][0][t] = tmp0;
      V_PS[ig][z][1][t] = tmp1;
      V_PS[ig][z][2][t] = tmp2;
    }
}

// solve policy function for industry i
int solve_policies_PS(int ig)
{
  double maxdiff = 999;
  int imaxdiff[3];

  // first do SS policy
  int iter=0;
  do
    {
      iter++;

      // case 1: firms know TPU ends with WTO accession
      if(perm_tpu==0)
	iterate_policies_PS(ig,NT-1,0,&maxdiff,imaxdiff);
      
      // case 2: firms think TPU will last forever
      else
	iterate_policies_PS(ig,NT-1,tpu_prob_ps,&maxdiff,imaxdiff);
    }
  while(maxdiff>policy_tol_abs && iter < policy_max_iter);
  
  if(iter==policy_max_iter)
    {
      printf("\tValue function iteration failed for industry %d! Diff = %0.4g\n",ig,maxdiff);
      return 1;
    }
  
  // now iterate backwards
  for(int t=NT-2; t>=0; t--)
    {
      // case 1: firms know TPU ends with WTO accession
      if(t>=t_wto && perm_tpu==0)
	iterate_policies_PS(ig,t,0,&maxdiff,imaxdiff);

      // case 2: firms think TPU will last forever
      else
	iterate_policies_PS(ig,t,tpu_prob_ps,&maxdiff,imaxdiff);
    }

  return 0;
}

// solve policies for all industries in parallel
int solve_policies2_PS()
{
  if(verbose)
    printf("\nSolving stochastic dynamic program with constant transition probabilities as in PS (2016)...\n");

  time_t start, stop;
  time(&start);

  int cnt=0;
  
#ifdef _OPENMP
#pragma omp parallel for
#endif
  for(int ig=0; ig<NG; ig++)
    {
      policy_solved_flag[ig] = solve_policies_PS(ig);
      cnt += policy_solved_flag[ig];
    }

  time(&stop);
  
  if(verbose)
    {
      printf("Finished dynamic programs in %0.0f seconds. %d failed to converge.\n",difftime(stop,start),cnt);
    }
  
  return 0;  
}


/////////////////////////////////////////////////////////////////////////////
// 5. Trade dynamics via distribution iteration
/////////////////////////////////////////////////////////////////////////////

const double dist_tol = 1.0e-11;
const int max_dist_iter = 5000;

double dist[NG][NZ][3] = {{{0.0}}};
double tmp_dist[NG][NZ][3] = {{{0.0}}};


double exports[NG][NT];
double quantity[NG][NT];
double pv_tau[NG][NT];
double nf[NG][NT];

// initialize distribution
void init_dist()
{
  for(int ig=0; ig<NG; ig++)
    {
      double sum=0.0;
      for(int iz=0; iz<NZ; iz++)
	{
	  tmp_dist[ig][iz][0] = 0.0;
	  tmp_dist[ig][iz][1] = 0.0;
	  tmp_dist[ig][iz][2] = 0.0;
	  
	  dist[ig][iz][0] = z_ucond_probs[good_2_ind[ig]][iz];
	  dist[ig][iz][1] = 0.0;
	  dist[ig][iz][2] = 0.0;
	  sum += dist[ig][iz][0];
	}
      if(fabs(sum-1.0)>1.0e-8)
	{
	  printf("\nInitial distribution does not sum to one! i = %d, sum = %0.4g\n",ig,sum);
	}

    }
}

// distribution iteration driver
int update_dist(int ig, int t, int reform_flag)
{
  int ii = good_2_ind[ig];
  
  for(int iz=0; iz<NZ; iz++)
    {
      for(int is=0; is<3; is++)
	{
	  tmp_dist[ig][iz][is]=0.0;
	}
    }

  for(int iz=0; iz<NZ; iz++)
    {
      double surv_prob = delta[ii][iz];
      
      for(int is=0; is<3; is++)
	{ 
	  int gex_ = 0;

	  // if we are before the 1980 reform, use the
	  // pre-MIT shock policy
	  if(t<t_reform)
	    {
	      if(ps==0)
		{ 
		  if(reform_flag==0)
		    gex_ = gex_ref_det[ig][iz][is][t];
		  else if(reform_flag==5)
		    gex_ = gex_ref_det[ig][iz][is][NT-1];
		  else if(reform_flag<3)
		    gex_ = gex_nntr[ig][iz][is];
		  else
		    gex_ = gex_markov[ig][iz][is][0][t];
		}
	      else
		{
		  gex_ = gex_nntr[ig][iz][is];
		}
	    }
	  
	  // otherwise it depends on which scenario we are using
	  else
	    {
	      if(ps==0)
		{
		  if(reform_flag==0 || reform_flag==1)
		    gex_ = gex_ref_det[ig][iz][is][t];
		  else if(reform_flag==5)
		    gex_ = gex_ref_det[ig][iz][is][NT-1];
		  else if(reform_flag == 2)
		    gex_ = gex_nntr[ig][iz][is];
		  else if(reform_flag==3)
		    gex_ = gex_markov[ig][iz][is][1][t];
		}
	      else
		{
		  if(ps==1)
		    {
		      if(t>=t_wto)
			gex_ = gex_ref_det[ig][iz][is][t];
		      else
			gex_ = gex_PS[ig][iz][is][t];		      
		    }
		  else
		    {
		      if(t>=t_wto || t<t_1990)
			gex_ = gex_ref_det[ig][iz][is][t];
		      else
			gex_ = gex_PS[ig][iz][is][t];

		    }
		}
	      
	    }

	  for(int izp=0; izp<NZ; izp++)
	    {
	      tmp_dist[ig][izp][0] += (1.0-surv_prob)*
		dist[ig][iz][is]*z_ucond_probs[ii][izp];

	      if(gex_==1)
		{
		  if(is==0)
		    {
		      tmp_dist[ig][izp][1] += dist[ig][iz][is]*
			surv_prob*z_trans_probs[ii][iz][izp];
		    }
		  else if(is==1)
		    {
		      tmp_dist[ig][izp][1] += dist[ig][iz][is]*
			surv_prob*z_trans_probs[ii][iz][izp]*rho0[ii];
		      
		      tmp_dist[ig][izp][2] += dist[ig][iz][is]*
			surv_prob*z_trans_probs[ii][iz][izp]*(1.0-rho0[ii]);
		    }
		  else if(is==2)
		    {
		      tmp_dist[ig][izp][1] += dist[ig][iz][is]*
			surv_prob*z_trans_probs[ii][iz][izp]*(1.0-rho1[ii]);
		      
		      tmp_dist[ig][izp][2] += dist[ig][iz][is]*
			surv_prob*z_trans_probs[ii][iz][izp]*rho1[ii];
		    }
		}
	      else
		{
		  tmp_dist[ig][izp][0] += surv_prob*
		    dist[ig][iz][is]*z_trans_probs[ii][iz][izp];
		}
	      
	    }
	}
    }

  double sum = 0.0;
  for(int iz=0; iz<NZ; iz++)
    {
      for(int is=0; is<3; is++)
	{
	  sum = sum+tmp_dist[ig][iz][is];
	}
    }

  if(fabs(sum-1.0)>1.0e-8)
    {
      printf("\nUpdated distribution does not sum to one! i = %d, t = %d, sum = %0.16f\n",ig,t,sum);
      return 1;
    }

  return 0;
}

void calc_trans_vars(int ig, int t, int reform_flag)
{
  int ii = good_2_ind[ig];
  
  double expart_rate=0.0;
  double total_exports=0.0;
  double total_quantity=0.0;

  double tau = tau_applied[ig][t];
  double tau_hat_ = pow(tau_applied[ig][t],-theta[ii]);
  double iceberg_hat_ = pow(iceberg_val[ig][t],1.0-theta[ii]);

  if(reform_flag==5)
    {
      tau = tau_applied[ig][NT-1];
      tau_hat_= pow(tau_applied[ig][NT-1],-theta[ii]);
      iceberg_hat_ = pow(iceberg_val[ig][NT-1],1.0-theta[ii]);
    }
  else if(t<t_reform || reform_flag==2)
    {
      tau = tau_nntr[ig];
      tau_hat_ = pow(tau_nntr[ig],-theta[ii]);
      iceberg_hat_ = pow(iceberg_val[ig][t],1.0-theta[ii]);
    }
  else if(t==t_reform && reform_flag != 2)
    {
      tau = frac_1980*tau_applied[ig][t] + (1.0-frac_1980)*tau_nntr[ig];
      tau_hat_ = (frac_1980*pow(tau_applied[ig][t],-theta[ii]) +
		  (1.0-frac_1980)*pow(tau_nntr[ig],-theta[ii]));
      iceberg_hat_ = pow(iceberg_val[ig][t],1.0-theta[ii]);
    }
  
  for(int z=0; z<NZ; z++)
    {
      for(int is=0; is<3; is++)
	{
	  double d=dist[ig][z][is];
	  double v=0;
	  double q=0;
	  double p = (theta[ii]/(theta[ii]-1.0)) * tau / z_grid[ii][z];
	  
	  int s=0;
	  if (is==1) // if it is currently a bad exporter...
	    {
	      v += theta_hat2[ii] * tau_hat_ * iceberg_hat_ *
		z_hat[ii][z]*pow(xi[ii],1.0-theta[ii]);
	      q += v/p;
	      s=1;
	    }
	  else if(is==2) // if it is current a good exporter...
	    {
	      v += theta_hat2[ii] * tau_hat_ * iceberg_hat_ * z_hat[ii][z];
	      q += v/p;
	      s=1;
	    }

	  total_exports += v * d;
	  total_quantity += q * d;
	  expart_rate += s * d;
	  
	}
    }

  exports[ig][t] = total_exports;
  quantity[ig][t] = total_quantity;
  nf[ig][t] = expart_rate;

  if(t>t_wto+30 && exports[ig][t]/exports[ig][t-1]<-0.05)
    {
      printf("Error!!\n");
    }

  
  return;
}

void calc_pv_tau(int reform_flag)
{
  for(int ig=0; ig<NG; ig++)
    {
      if(reform_flag<2 || reform_flag==5)
	{
	  pv_tau[ig][NT-1] = tau_applied[ig][NT-1]/(1.0-Q);
	
	  for(int t=NT-2; t>=0; t--)
	    {
	      if(t>=t_reform)
		{
		  pv_tau[ig][t] = tau_applied[ig][t] + Q*pv_tau[ig][t+1];
		}
	      else
		{
		  if(reform_flag==0)
		    {
		      pv_tau[ig][t] = tau_nntr[ig] + Q*pv_tau[ig][t+1];
		    }
		  else
		    {
		      pv_tau[ig][t] = tau_nntr[ig]/(1.0-Q);
		    }
		}
	    }
	}
      else if(reform_flag==2)
	{
	  for(int t=0; t<NT; t++)
	    {
	      pv_tau[ig][t] = tau_nntr[ig]/(1.0-Q);
	    }
	}
      else if(reform_flag==3 && ps==0)
	{
	  if(perm_tpu==0)
	    {
	      // use markov chain to figure out terminal PVs
	      double tau_[2];
	      tau_[0] = tau_nntr[ig];
	      tau_[1] = tau_applied[ig][NT-1];

	      double trans_mat_[2][2];
	      trans_mat_[0][0] = tpu_trans_mat[0][0][NT-1];
	      trans_mat_[0][1] = tpu_trans_mat[0][1][NT-1];
	      trans_mat_[1][0] = tpu_trans_mat[1][0][NT-1];
	      trans_mat_[1][1] = tpu_trans_mat[1][1][NT-1];
	  
	      markov_stationary_dist(2, trans_mat_, tau_);
	      tau_[0] = tau_[0]/(1.0-Q);
	      tau_[1] = tau_[1]/(1.0-Q);
	      pv_tau[ig][NT-1] = tau_[1];

	      for(int t=NT-2; t>=0; t--)
		{
		  for(int p=0; p<2; p++)
		    {
		      double tmp0=0;
		      double tmp1=0;
		      if(t>=t_reform)
			{
			  tmp0 = tau_nntr[ig] + Q*(tpu_trans_mat[0][0][t]*tau_[0] + tpu_trans_mat[0][1][t]*tau_[1]);
			  tmp1 = tau_applied[ig][t] + Q*(tpu_trans_mat[1][0][t]*tau_[0] + tpu_trans_mat[1][1][t]*tau_[1]);
			  pv_tau[ig][t] = tmp1;
			}
		      else
			{
			  tmp0 = tau_nntr[ig] + Q*(tpu_trans_mat[0][0][t]*tau_[0] + tpu_trans_mat[0][1][t]*tau_[1]);
			  tmp1 = tau_applied[ig][t_reform] + Q*(tpu_trans_mat[1][0][t]*tau_[0] + tpu_trans_mat[1][1][t]*tau_[1]);
			  pv_tau[ig][t] = tmp0;
			}

		      tau_[0] = tmp0;
		      tau_[1] = tmp1;
		    }
		}
	    }
	  else
	    {
	      for(int t_outer=0; t_outer<NT; t_outer++)
		{
		  double tau_[2];
		  tau_[0] = tau_nntr[ig];
		  tau_[1] = tau_applied[ig][NT-1];

		  double trans_mat_[2][2];
		  trans_mat_[0][0] = tpu_trans_mat[0][0][NT-1];
		  trans_mat_[0][1] = tpu_trans_mat[0][1][NT-1];
		  trans_mat_[1][0] = tpu_trans_mat[1][0][NT-1];
		  trans_mat_[1][1] = tpu_trans_mat[1][1][NT-1];
	      
		  markov_stationary_dist(2, trans_mat_, tau_);
		  tau_[0] = tau_[0]/(1.0-Q);
		  tau_[1] = tau_[1]/(1.0-Q);
		  pv_tau[ig][NT-1] = tau_[1];

		  for(int t=NT-2; t>=t_outer; t--)
		    {
		      for(int p=0; p<2; p++)
			{
			  double tmp0=0;
			  double tmp1=0;
			  if(t>=t_reform)
			    {
			      tmp0 = tau_nntr[ig] + Q*(tpu_trans_mat[0][0][t]*tau_[0] + tpu_trans_mat[0][1][t]*tau_[1]);
			      tmp1 = tau_applied[ig][t] + Q*(tpu_trans_mat[1][0][t]*tau_[0] + tpu_trans_mat[1][1][t]*tau_[1]);
			      pv_tau[ig][t] = tmp1;
			    }
			  else
			    {
			      tmp0 = tau_nntr[ig] + Q*(tpu_trans_mat[0][0][t]*tau_[0] + tpu_trans_mat[0][1][t]*tau_[1]);
			      tmp1 = tau_applied[ig][t_reform] + Q*(tpu_trans_mat[1][0][t]*tau_[0] + tpu_trans_mat[1][1][t]*tau_[1]);
			      pv_tau[ig][t] = tmp0;
			    }

			  tau_[0] = tmp0;
			  tau_[1] = tmp1;
			}
		    }
		}
	    }
		    
	}

      else if(reform_flag==3 && ps)
	{
	  if(perm_tpu==1)
	    {
	      double pv_nntr = tau_nntr[ig]/(1.0-Q);
	      double pv_notpu = tau_applied[ig][NT-1]/(1.0-Q);
	      pv_tau[ig][NT-1] = (tau_applied[ig][NT-1] + tpu_prob_ps*Q*pv_nntr)/(1.0-Q*(1.0-tpu_prob_ps));

	      
	      for(int t=NT-2; t>=0; t--)
		{
		  pv_notpu = tau_applied[ig][t] + Q*pv_notpu;

		  if(t<t_reform)
		    {
		      pv_tau[ig][t] = pv_nntr;
		    }
		  else if(t<t_wto && (ps==1 || t>=t_1990))
		    {
		      pv_tau[ig][t] = tau_applied[ig][t] + Q*(tpu_prob_ps*pv_nntr + (1.0-tpu_prob_ps)*pv_tau[ig][t+1]);
		    }
		  else
		    {
		      pv_tau[ig][t] = pv_notpu;
		    }
		}
	    }
	  else
	    {
	      double pv_nntr = tau_nntr[ig]/(1.0-Q);
	      double pv_notpu = tau_applied[ig][NT-1]/(1.0-Q);
	      pv_tau[ig][NT-1] = pv_notpu;
	      
	      for(int t=NT-2; t>=0; t--)
		{
		  pv_notpu = tau_applied[ig][t] + Q*pv_notpu;

		  if(t>=t_wto)
		    {
		      pv_tau[ig][t] = pv_notpu;
		    }
		  else if(t>=t_reform)
		    {
		      if(ps==1)
			{
			  pv_tau[ig][t] = tau_applied[ig][t] + Q*(tpu_prob_ps*pv_nntr + (1.0-tpu_prob_ps)*pv_tau[ig][t+1]);
			}
		      else
			{
			  if(t>=t_1990)
			    {
			      pv_tau[ig][t] = tau_applied[ig][t] + Q*(tpu_prob_ps*pv_nntr + (1.0-tpu_prob_ps)*pv_tau[ig][t+1]);
			    }
			  else
			    {
			      pv_tau[ig][t] = pv_notpu;
			    }
			}
		    }
		  else
		    {
		      pv_tau[ig][t] = pv_nntr;
		    }
		  
		}

	    }
		    
	}

      
      // annualize it
      for(int t=0; t<NT; t++)
	{
	  pv_tau[ig][t] = pv_tau[ig][t]*(1.0-Q);
	}
      
    }
}

void do_trans_dyn(int reform_flag)
{
  printf("\nComputing transition dynamics for scenario %d...\n",reform_flag);

  time_t start, stop;
  time(&start);

  if(reform_flag != 4)
    init_dist();

  int reform_flag_ = reform_flag;
  if(reform_flag==4)
    reform_flag_ = 1;
  
  for(int t = 0; t<NT; t++)
    {
      
#ifdef _OPENMP
#pragma omp parallel for
#endif
      for(int ig=0; ig<NG; ig++)
	{
	  calc_trans_vars(ig,t, reform_flag_);
	  update_dist(ig,t, reform_flag_);
	}
      
      memcpy(dist,tmp_dist,sizeof(double)*(NG*NZ*3));
    }


  calc_pv_tau(reform_flag);
  
  char suffix[24] = "baseline";

  if(one_sector==1)
    {
      snprintf(suffix,24,"one_sector");
    }   
  else if(one_sector==2)
    {
      snprintf(suffix,24,"avg_sector");
    }
  else if(ps==1)
    {
      if(perm_tpu==0)
	snprintf(suffix,24,"PS_1980_temptpu");
      else
	snprintf(suffix,24,"PS_1980_permtpu");
    }
  else if(ps==2)
    {
      if(perm_tpu==0)
	snprintf(suffix,24,"PS_1990_temptpu");
      else
	snprintf(suffix,24,"PS_1990_permtpu");
    }
  else if(sens_xi==1)
    {
      snprintf(suffix,24,"lo_xi");
    }
  else if(sens_xi==2)
    {
      snprintf(suffix,24,"hi_xi");
    }
  else if(sens_xi==3)
    {
      snprintf(suffix,24,"lo_rhoxi");
    }
  else if(sens_xi==4)
    {
      snprintf(suffix,24,"hi_rhoxi");
    }
  else if(alt_coeffs==0)
    {
      if(sunk_cost==1 && calvo==0 && perm_tpu==0)
	snprintf(suffix,24,"sunkcost");
      else if(calvo==1 && sunk_cost==0 && perm_tpu==0)
	snprintf(suffix,24,"calvo");
      else if(calvo==0 && sunk_cost==0 && perm_tpu==1)
	snprintf(suffix,24,"permtpu");
    }
  else if(alt_coeffs)
    {
      if(alt_coeffs==5)
	snprintf(suffix,24,"ci_lower");
      else if(alt_coeffs==6)
	snprintf(suffix,24,"ci_upper");
    }

  char fname[128] = "";
  if(reform_flag==0)
    {
      snprintf(fname,128,OUTPUT "simul_agg_det0_%s.csv",suffix);
    }
  else if(reform_flag==1)
    {
      snprintf(fname,128,OUTPUT "simul_agg_det1_%s.csv",suffix);
    }
  else if(reform_flag==2)
    {
      snprintf(fname,128,OUTPUT "simul_agg_det2_%s.csv",suffix);
    }
  else if(reform_flag==4)
    {
      snprintf(fname,128,OUTPUT "simul_agg_det4_%s.csv",suffix);
    }
  else if(reform_flag==5)
    {
      snprintf(fname,128,OUTPUT "simul_agg_det5_%s.csv",suffix);
    }
  else if(reform_flag==3)
    {
      snprintf(fname,128,OUTPUT "simul_agg_tpu_%s.csv",suffix);
    }

  FILE * file2 = fopen(fname,"w");
  fprintf(file2,"i,s,y,tau_applied,tau_nntr,iceberg,gap,exports,quantity,num_exporters,pv_tau\n");

  for(int ig=0; ig<NG; ig++)
    {
      for(int t=1; t<NT; t++)
	{	  
	  fprintf(file2,"%d,%d,%d,%0.16f,%0.16f,%0.16f,%0.16f,%0.16f,%0.16f,%0.16f,%0.16f\n",
		  ig,good_2_ind[ig],t,tau_applied[ig][t],tau_nntr[ig],iceberg_val[ig][t],gap[ig],
		  exports[ig][t],quantity[ig][t],nf[ig][t],pv_tau[ig][t]);
	}
    }
  
  fclose(file2);

  time(&stop);

  if(verbose)
    printf("Transitions complete in %0.0f seconds.\n",difftime(stop,start));

}

int calc_cal_moments(int ii, double *expart_rate_, double *exit_rate_, double *inc_prem_, double *size_dist_)
{ 
  *expart_rate_ = 0;
  *exit_rate_ = 0;
  *inc_prem_ = 0;
  *size_dist_ = 0;  

  double total_mass = 0;
  
  double avg_size = 0;
  double exporter_mass=0;
  
  double new_size = 0;
  double new_mass = 0;
  
  double inc_size=0;
  double inc_mass=0;

  // export participation rate, exit rate, incumbent premium
  for(int ig=0; ig<NG; ig++)
    {
      if( (one_sector==1 && ii==NI) || good_2_ind[ig]==ii)
      //if(good_2_ind[ig]==ii)
	{
	  int ii_ = good_2_ind[ig];
	  
	  for(int z=0; z<NZ; z++)
	    {
	      double Ezh = 0;
	      for(int zp=0; zp<NZ; zp++)
		{
		  Ezh += z_trans_probs[ii_][z][zp] * z_hat[ii_][zp];
		}
      	  
	      double tau_hat_ = pow(tau_applied[ig][NT-1],-theta[ii_]);
	      double iceberg_hat_ = 1.0;
	  
	      for(int e=0; e<3; e++)
		{
		  double d = dist[ig][z][e];
		  total_mass += d;

		  //int gex_ = gex_ref_det[ig][z][e][NT-1];
		  int gex_ = gex_markov[ig][z][e][1][NT-1];
	  
		  if(e>0)
		    {
		      exporter_mass += d;
		      
		      if(gex_==0)
			*exit_rate_ += d;
		      else
			*exit_rate_ += (1.0-delta[ii_][z])*d;
	      
		      if(e==1)
			avg_size += theta_hat2[ii_]*tau_hat_*iceberg_hat_*z_hat[ii_][z]*
			  pow(xi[ii_],1.0-theta[ii_])*d;
		      else
			avg_size += theta_hat2[ii_]*tau_hat_*iceberg_hat_*z_hat[ii_][z]*d;
		    }
	  
		  if(e==0 && gex_==1)
		    {
		      new_size += theta_hat2[ii_]*tau_hat_*iceberg_hat_*Ezh*
			pow(xi[ii_],1.0-theta[ii_])*d*delta[ii_][z];
		      new_mass += d*delta[ii_][z];
		    }

		  if(e>0 && gex_==1)
		    {
		      double y1 = theta_hat2[ii_]*tau_hat_*iceberg_hat_*Ezh*pow(xi[ii_],1.0-theta[ii_])*d*delta[ii_][z];
		      double y2 = theta_hat2[ii_]*tau_hat_*iceberg_hat_*Ezh*d*delta[ii_][z];
		      if(e==1)
			{
			  inc_size += rho0[ii_]*y1 + (1.0-rho0[ii_])*y2;
			}
		      else
			{
			  inc_size += rho1[ii_]*y2 + (1.0-rho1[ii_])*y1;
			}
		      inc_mass += d*delta[ii_][z];
		    } 
		}
	    }
	}
    }

  *expart_rate_ = exporter_mass/total_mass;
  avg_size = avg_size / exporter_mass;  
  *exit_rate_ = *exit_rate_ / exporter_mass;
  *inc_prem_ = (inc_size/inc_mass) / (new_size/new_mass);

  // no go back and do the size distribution measure, now that we know the average exporter size
  // export participation rate, exit rate, incumbent premium
  double variance = 0.0;
  for(int ig=0; ig<NG; ig++)
    {
      if( (one_sector==1 && ii==NI) || good_2_ind[ig]==ii)
	{
	  int ii_ = good_2_ind[ig];
	  
	  for(int z=0; z<NZ; z++)
	    {      	  
	      double tau_hat_ = pow(tau_applied[ig][NT-1],-theta[ii_]);
	      double iceberg_hat_ = 1.0;
	  
	      for(int e=0; e<3; e++)
		{
		  double d = dist[ig][z][e];

		  double y = 0.0;
		  if(e>0)
		    {	      
		      if(e==1)
			y = theta_hat2[ii_]*tau_hat_*iceberg_hat_*z_hat[ii_][z]*
			  pow(xi[ii_],1.0-theta[ii_]);
		      else
			y= theta_hat2[ii_]*tau_hat_*iceberg_hat_*z_hat[ii_][z];

		      variance += fabs((y-avg_size)*(y-avg_size)) * d;
		    }
	  
		}
	    }
	}
    }

  variance = variance/exporter_mass;
  *size_dist_ = log(sqrt(variance)/avg_size);

  return 0;
}



/////////////////////////////////////////////////////////////////////////////
// 5a. Calibrating exporter parameters in long run steady state
/////////////////////////////////////////////////////////////////////////////

void save_exporter_params()
{
  FILE * file = 0x0;
  if(one_sector==0 && multi_sector==1 && calvo==0 && sunk_cost==0)
    file = fopen(OUTPUT "exporter_params_baseline.csv","w");
  else if(one_sector==1)
    file = fopen(OUTPUT "exporter_params_one_sector.csv","w");
  else if(one_sector==2)
    file = fopen(OUTPUT "exporter_params_avg_sector.csv","w");
  else if(calvo)
    file = fopen(OUTPUT "exporter_params_calvo.csv","w");
  else if(sunk_cost)
    file = fopen(OUTPUT "exporter_params_sunkcost.csv","w");
  
  for(int ii=0; ii<NI; ii++)
    {
      fprintf(file,"%d %0.16f %0.16f %0.16f %0.16f\n",ii,kappa0[ii],kappa1[ii],xi[ii],sig_z[ii]);
    }
  fclose(file);
}

int load_exporter_params()
{
  FILE * file = 0x0;
  if(one_sector==0 && multi_sector==1 && calvo==0 && sunk_cost==0)
    file = fopen(OUTPUT "exporter_params_baseline.csv","r");
  else if(one_sector==1)
    file = fopen(OUTPUT "exporter_params_one_sector.csv","r");
  else if(one_sector==2)
    file = fopen(OUTPUT "exporter_params_avg_sector.csv","r");
  else if(calvo)
    file = fopen(OUTPUT "exporter_params_calvo.csv","r");
  else if(sunk_cost)
    file = fopen(OUTPUT "exporter_params_sunkcost.csv","r");

  if(file == NULL)
    {
      printf("Failed to open file with exporter parameters!\n");
      return 1;
    }

  int ii;
  double kappa0_, kappa1_, xi_, sig_z_;
  int cnt=0;  
  while(cnt<NI && fscanf(file,"%d %lf %lf %lf %lf",&ii ,&kappa0_, &kappa1_, &xi_, &sig_z_) == 5)
    {
      cnt+=1;
      
      if(ii<0 || ii>NI-1 || gsl_isnan(kappa0_) || gsl_isnan(kappa1_) ||
	 gsl_isnan(xi_) || gsl_isinf(sig_z_))
	{
	  printf("Bad value in exporter params!\n");
	  fclose(file);
	  return 1;
	}

      kappa0[ii]=kappa0_;
      kappa1[ii]=kappa1_;
      xi[ii]=xi_;
      sig_z[ii]=sig_z_;
    }
  fclose(file);

  if(cnt!=NI)
    {
      printf("Failed to load exporter params!\n");
      return 1;
    }

  return 0;
}

int calibrate_exporter_params()
{
  printf("Calibrating exporter parameters...\n");

  time_t start, stop;
  time(&start);

  double maxdiff=0.0;
  int iter=0;
  do
    {
      iter++;
      linebreak2();
      save_exporter_params();

      // -------------------------------------------------------------------------------------
      // reset constants and re-discretize productivity processes
      for(int ii=0; ii<NI; ii++)
	{
	  theta_hat[ii] = (1.0/theta[ii]) * pow(theta[ii]/(theta[ii]-1.0),1.0-theta[ii]);
	  theta_hat2[ii] = theta[ii] * theta_hat[ii];
	  discretize_z(ii);
	  calc_death_probs(ii);
	}

      // -------------------------------------------------------------------------------------
      // compute long-run exporter dynamics moments by sector
      
      // solve for transition dynamics
      if(solve_policies2_markov())
	return 1;
      
      do_trans_dyn(3);

      // compute moments and update params
      maxdiff = 0.0;
      if(one_sector)
	{
	  printf("\nParams 1  = (%0.4f %0.4f %0.4f %0.4f)\n",kappa0[0],kappa1[0],xi[0],sig_z[0]);
	  
	  double expart_rate_=0.0;
	  double exit_rate_=0.0;
	  double inc_prem_=0.0;
	  double size_dist_=0;
	  int ii_=0;

	  if(one_sector==1)
	    {
	      ii_ = NI;
	      calc_cal_moments(ii_,&expart_rate_,&exit_rate_,&inc_prem_,&size_dist_);
	    }
	  else if(one_sector==2)
	    {
	      ii_ = NI-1;
	      
	      for(int ii=0; ii<NI; ii++)
		{
		  double expart_rate2_=0.0;
		  double exit_rate2_=0.0;
		  double inc_prem2_=0.0;
		  double size_dist2_=0;
		  
		  calc_cal_moments(ii_,&expart_rate2_,&exit_rate2_,&inc_prem2_,&size_dist2_);
		  
		  expart_rate_ += expart_rate2_;
		  exit_rate_ += exit_rate2_;
		  inc_prem_ += inc_prem2_;
		  size_dist_ += size_dist2_;		
		}

	      expart_rate_ /= NI;
	      exit_rate_ /= NI;
	      inc_prem_ /= NI;
	      size_dist_ /= NI;
	    }

	  printf("Export participation rate = %0.4f (%0.4f)\n",expart_rate_, export_participation[ii_]);
	  printf("Exit rate rate = %0.4f (%0.4f)\n",exit_rate_,exit_rate[ii_]);
	  printf("Incumbent premium = %0.4f (%0.4f)\n",inc_prem_,incumbent_premium[ii_]);
	  printf("Log CV of exports = %0.4f (%0.4f)\n",size_dist_,size_dist[ii_]);

	  double tmp1 = (expart_rate_ - export_participation[ii_])/export_participation[ii_];
	  double tmp2 = (exit_rate_ - exit_rate[ii_])/exit_rate[ii_];
	  double tmp3 = (inc_prem_ - incumbent_premium[ii_])/incumbent_premium[ii_];
	  double tmp4 = (size_dist_ - size_dist[ii_])/size_dist[ii_];

	  for(int ii=0; ii<NI; ii++)
	    {
	      if(calvo==0)
		{
		  kappa0[ii] *= (1.0 + sign(tmp1) * log(1.0 + fabs(tmp1)) * cal_expart_mult);
		  kappa1[ii] *= (1.0 - sign(tmp2) * log(1.0 + fabs(tmp2)) * cal_exit_mult);
		  xi[ii] *= (1.0 - sign(tmp3) * log(1.0 + fabs(tmp3)) * cal_inc_prem_mult);
		  sig_z[ii] *= (1.0 - sign(tmp4) * log(1.0 + fabs(tmp4)) * cal_size_dist_mult);
		}
	      else
		{
		  tmp3=0.0;		  
		  kappa0[ii] *= (1.0 + sign(tmp1) * log(1.0 + fabs(tmp1)) * cal_expart_mult);
		  kappa1[ii] = 0.0;
		  xi[ii] = 1.0;
		  sig_z[ii] *= (1.0 - sign(tmp4) * log(1.0 + fabs(tmp4)) * cal_size_dist_mult);
		  delta1[ii] *= (1.0 - sign(tmp2) * log(1.0 + fabs(tmp2)) * cal_exit_mult);
		}
	    }
	  
	  if(fabs(tmp1)>maxdiff)
	    maxdiff=fabs(tmp1);

	  if(fabs(tmp2)>maxdiff)
	    maxdiff=fabs(tmp2);

	  if(fabs(tmp3)>maxdiff)
	    maxdiff = fabs(tmp3);

	  if(fabs(tmp4)>maxdiff)
	    maxdiff=fabs(tmp4);
	}
      else
	{
	  for(int ii=0; ii<NI; ii++)
	    {
	      printf("\n\tIndustry %d\n",ii);
	      printf("\tParams  = (%0.4f %0.4f %0.4f %0.4f)\n",kappa0[ii],kappa1[ii],xi[ii],sig_z[ii]);
	      
	      double expart_rate_=0.0;
	      double exit_rate_=0.0;
	      double inc_prem_=0.0;
	      double size_dist_=0;
	      
	      calc_cal_moments(ii, &expart_rate_,&exit_rate_,&inc_prem_,&size_dist_);

	      printf("\tExport participation rate = %0.4f (%0.4f)\n",expart_rate_, export_participation[ii]);
	      printf("\tExit rate rate = %0.4f (%0.4f)\n",exit_rate_,exit_rate[ii]);
	      printf("\tIncumbent premium = %0.4f (%0.4f)\n",inc_prem_,incumbent_premium[ii]);
	      printf("\tLog CV of exports = %0.4f (%0.4f)\n",size_dist_,size_dist[ii]);

	      double tmp1 = (expart_rate_ - export_participation[ii])/export_participation[ii];
	      double tmp2 = (exit_rate_ - exit_rate[ii])/exit_rate[ii];
	      double tmp3 = (inc_prem_ - incumbent_premium[ii])/incumbent_premium[ii];
	      double tmp4 = (size_dist_ - size_dist[ii])/size_dist[ii];

	      if(calvo==0)
		{
		  kappa0[ii] *= (1.0 + sign(tmp1) * log(1.0 + fabs(tmp1)) * cal_expart_mult);
		  kappa1[ii] *= (1.0 - sign(tmp2) * log(1.0 + fabs(tmp2)) * cal_exit_mult);
		  xi[ii] *= (1.0 - sign(tmp3) * log(1.0 + fabs(tmp3)) * cal_inc_prem_mult);
		  sig_z[ii] *= (1.0 - sign(tmp4) * log(1.0 + fabs(tmp4)) * cal_size_dist_mult);
		}
	      else
		{
		  tmp3=0.0;		  
		  kappa0[ii] *= (1.0 + sign(tmp1) * log(1.0 + fabs(tmp1)) * cal_expart_mult);
		  kappa1[ii] = 0.0;
		  xi[ii] = 1.0;
		  sig_z[ii] *= (1.0 - sign(tmp4) * log(1.0 + fabs(tmp4)) * cal_size_dist_mult);
		  delta1[ii] *= (1.0 - sign(tmp2) * log(1.0 + fabs(tmp2)) * cal_exit_mult);
 		}
	      
	      if(fabs(tmp1)>maxdiff)
		maxdiff=fabs(tmp1);
	      
	      if(fabs(tmp2)>maxdiff)
		maxdiff=fabs(tmp2);
	      
	      if(fabs(tmp3)>maxdiff)
		maxdiff = fabs(tmp3);
	      
	      if(fabs(tmp4)>maxdiff)
		maxdiff = fabs(tmp4);
	    }
	}
      
      printf("\nIter = %d, error = %0.6f\n",iter, maxdiff);
    }
  while(iter<max_cal_iter && maxdiff>cal_tol);

  /*
  if(calvo==0)
    {
      printf("\nParams 2  = (%0.4f %0.4f)\n",frac_1980,rho0[0]);
  
      if(solve_policies2_det())
	return 1;
  
      do_trans_dyn(1);

      char buffer[128];
      if(one_sector==1)
	{
	  sprintf(buffer,"python3 -W ignore ../scripts/proc_simul.py -one-sector -calc-te");	  
	}
      else if(one_sector==2)
	{
	  sprintf(buffer,"python3 -W ignore ../scripts/proc_simul.py -avg-sector -calc-te");
	}
      else
	{
	  sprintf(buffer,"python3 -W ignore ../scripts/proc_simul.py -calc-te");
	}

        if(system(buffer))
	return 1;
	}*/

  time(&stop);
  printf("\nAnalysis complete! Runtime = %0.0f seconds.\n",
	 difftime(stop,start));
 
  return 0;

}

/////////////////////////////////////////////////////////////////////////////
// 5b. Calibrating TPU probs
/////////////////////////////////////////////////////////////////////////////
double caldata[NT]={0.0};
double caldata_PS = 0.0;
int iter=0;

int load_caldata()
{
  FILE * file = 0x0;

  if(one_sector==1)
    {
      file = fopen(INPUT "caldata_one_sector.txt","r");
    }
  else if(one_sector==2)
    {
      file = fopen(INPUT "caldata_avg_sector.txt","r");
    }
  else if(ps==1)
    {
      if(perm_tpu)
	file = fopen(INPUT "caldata_PS_1980_permtpu.txt","r");
      else
	file = fopen(INPUT "caldata_PS_1980_temptpu.txt","r");
    }
  else if(ps==2)
    {
      if(perm_tpu)
	file = fopen(INPUT "caldata_PS_1990_permtpu.txt","r");
      else
	file = fopen(INPUT "caldata_PS_1990_temptpu.txt","r");
    }
  else if(sens_xi==1)
    {
      file = fopen(INPUT "caldata_lo_xi.txt","r");
    }
  else if(sens_xi==2)
    {
      file = fopen(INPUT "caldata_hi_xi.txt","r");
    }
  else if(sens_xi==3)
    {
      file = fopen(INPUT "caldata_lo_rhoxi.txt","r");
    }
  else if(sens_xi==4)
    {
      file = fopen(INPUT "caldata_hi_rhoxi.txt","r");
    }
  else if(calvo==1)
    {
      file = fopen(INPUT "caldata_calvo.txt","r");
    }
  else if(alt_coeffs==0)
    {
      if(perm_tpu==1)
	file = fopen(INPUT "caldata_permtpu.txt","r");
      else
	file = fopen(INPUT "caldata_baseline.txt","r");
    }
  else if(alt_coeffs==5)
    file = fopen(INPUT "caldata_ci_lower.txt","r");
  else if(alt_coeffs==6)
    file = fopen(INPUT "caldata_ci_upper.txt","r");
    
  if(!file)
    {
      printf("Failed to open file with calibration data!\n");
      return 1;
    }
  else
    {
      if(ps==0)
	{
	  int got = 0;
	  for(int t=3; t<t_data_max; t++)
	    {
	      got += fscanf(file,"%lf",&(caldata[t]));
	    }
	  fclose(file);
	  if(got != t_data_max-3)
	    {
	      printf("Failed to load calibration data!\n");
	      return 1;
	    }
	  else
	    {
	      return 0;
	    }
	}
      else
	{
	  int got = fscanf(file,"%lf",&caldata_PS);
	  if(got!=1)
	    {
	      printf("Failed to load calibration data!\n");
	      return 1;
	    }
	  else
	    {
	      return 0;
	    }

	}
    }
  
}

double calc_coeffs()
{
  linebreak2();
  printf("Computing differences between actual and simulated NNTR gap coefficients...\n\n");

  if(ps)
    {
      printf("P(MFN-->NNTR) = %0.6f\n",tpu_prob_ps);
    }
  else
    {
      if(alt_coeffs==2)
	{
	  printf("P(NNTR-->MFN) =");
	  for(int t=2; t<t_reform; t++)
	    {
	      printf(" %0.2f",tpu_trans_mat[0][1][t]);
	    }
	  printf("\n");
	}
      else
	{
	  printf("P(NNTR-->MFN) = %0.3f\n",tpu_trans_mat[0][1][0]);
	}
  
      printf("\nP(MFN-->NNTR) = \n");
      for(int t=t_reform; t<t_data_max-1; t++)
	{
	  printf("%d %0.2f\n",t+1973,tpu_trans_mat[1][0][t]);
	}
      printf("\n");
    }

  int baseline=0;
  FILE * file = 0x0;

  if(one_sector==1)
    {
      file = fopen(OUTPUT "tpuprobs_one_sector.txt","w");
    }
  else if(one_sector==2)
    {
      file = fopen(OUTPUT "tpuprobs_avg_sector.txt","w");
    }
  else if(ps==1)
    {
      if(perm_tpu==1)
	file = fopen(OUTPUT "tpuprobs_PS_1980_permtpu.txt","w");
      else
	file = fopen(OUTPUT "tpuprobs_PS_1980_temptpu.txt","w");
    }
  else if(ps==2)
    {
      if(perm_tpu==1)
	file = fopen(OUTPUT "tpuprobs_PS_1990_permtpu.txt","w");
      else
	file = fopen(OUTPUT "tpuprobs_PS_1990_temptpu.txt","w");
    }
  else if(calvo==1)
    {
      file = fopen(OUTPUT "tpuprobs_calvo.txt","w");
    }
  else if(sens_xi==1)
    {
      file = fopen(OUTPUT "tpuprobs_lo_xi.txt","w");
    }
  else if(sens_xi==2)
    {
      file = fopen(OUTPUT "tpuprobs_hi_xi.txt","w");
    }
  else if(sens_xi==3)
    {
      file = fopen(OUTPUT "tpuprobs_lo_rhoxi.txt","w");
    }
  else if(sens_xi==4)
    {
      file = fopen(OUTPUT "tpuprobs_hi_rhoxi.txt","w");
    }
  else if(alt_coeffs==0)
    {
      if(perm_tpu==1)
	file = fopen(OUTPUT "tpuprobs_permtpu.txt","w");
      else
	{
	  file = fopen(OUTPUT "tpuprobs_baseline.txt","w");
	  baseline=1;
	  printf("rho_xi = %0.4f\n",rho0[0]);
	}
    }
  else if(alt_coeffs==1)
    file = fopen(OUTPUT "tpuprobs_tsusa.txt","w");
  else if(alt_coeffs==2)
    file = fopen(OUTPUT "tpuprobs_inv.txt","w");
  else if(alt_coeffs==3)
    file = fopen(OUTPUT "tpuprobs_sitc68.txt","w");
  else if(alt_coeffs==4)
    file = fopen(OUTPUT "tpuprobs_sitc7.txt","w");
  else if(alt_coeffs==5)
    file = fopen(OUTPUT "tpuprobs_ci_lower.txt","w");
  else if(alt_coeffs==6)
    file = fopen(OUTPUT "tpuprobs_ci_upper.txt","w");

  if(ps==0)
    {
      int cnt=0;
      for(int t=2; t<t_reform; t++)
	{
	  cnt++;
	  fprintf(file,"%0.16f ",tpu_trans_mat[0][1][t]);
	}
      for(int t=t_reform; t<t_data_max; t++)
	{
	  cnt++;
	  fprintf(file,"%0.16f ",tpu_trans_mat[1][0][t]);
	}
      if(baseline==1)
	fprintf(file,"%0.16f",rho0[0]);
    }
  else
    {
      fprintf(file,"%0.16f",tpu_prob_ps);
    }
  fclose(file);

  time_t start, stop;
  time(&start);

  if(ps)
    {
      if(solve_policies2_PS())
	return -99;
    }
  else
    {
      if(solve_policies2_markov())
	return -99;
    }
  
  do_trans_dyn(3);
  
  time_t start2, stop2;
  time(&start2);
  printf("\nProcessing simulated data...\n");
  char buffer[128];

  if(one_sector==1)
    {
      sprintf(buffer,"python3 -W ignore ../scripts/proc_simul.py -one-sector");
    }
  else if(one_sector==2)
    {
      sprintf(buffer,"python3 -W ignore ../scripts/proc_simul.py -avg-sector");
    }
  else if(ps==1)
    {
      if(perm_tpu==1)
	sprintf(buffer,"python3 -W ignore ../scripts/proc_simul_PS.py -1980 -permtpu");
      else
	sprintf(buffer,"python3 -W ignore ../scripts/proc_simul_PS.py -1980 -temptpu");
    }
  else if(ps==2)
    {
      if(perm_tpu==1)
	sprintf(buffer,"python3 -W ignore ../scripts/proc_simul_PS.py -1990 -permtpu");
      else
	sprintf(buffer,"python3 -W ignore ../scripts/proc_simul_PS.py -1990 -temptpu");
    }
  else if(calvo==1)
    {
      sprintf(buffer,"python3 -W ignore ../scripts/proc_simul.py -calvo");
    }
  else if(sens_xi==1)
    {
      sprintf(buffer,"python3 -W ignore ../scripts/proc_simul.py -lo-xi");
    }    
  else if(sens_xi==2)
    {
      sprintf(buffer,"python3 -W ignore ../scripts/proc_simul.py -hi-xi");
    }    
  else if(sens_xi==3)
    {
      sprintf(buffer,"python3 -W ignore ../scripts/proc_simul.py -lo-rhoxi");
    }    
  else if(sens_xi==4)
    {
      sprintf(buffer,"python3 -W ignore ../scripts/proc_simul.py -hi-rhoxi");
    }    
  else if(alt_coeffs==0)
    {
      if(perm_tpu==1)
	sprintf(buffer,"python3 -W ignore ../scripts/proc_simul.py -permtpu");
      else
	sprintf(buffer,"python3 -W ignore ../scripts/proc_simul.py");
    }
  else
    {
      if(alt_coeffs==5)
	sprintf(buffer,"python3 -W ignore ../scripts/proc_simul.py -ci-lower");
      else if(alt_coeffs==6)
	sprintf(buffer,"python3 -W ignore ../scripts/proc_simul.py -ci-upper");

    }
  
  if(system(buffer))
    return -99;
  time(&stop2);
  printf("Processing complete in %0.0f seconds\n",difftime(stop2,start2));  

  if(load_caldata())
    return -99;

  double retval = -HUGE_VAL;
  if(ps)
    {
      printf("\nError = %0.2f\n",caldata_PS);
      retval = fabs(caldata_PS);
    }
  else
    {
      printf("\nErrors:\n");
      retval = -HUGE_VAL;
      double avg_pre80 = 0.0;
      for(int t=3; t<=t_reform; t++)
	{
	  avg_pre80 += fabs(caldata[t])*fabs(caldata[t]);
	}
      avg_pre80 = sqrt(avg_pre80/((double)(t_reform-3+1)));
      printf("Avg 1974-1979 = %0.2f\n\n",avg_pre80);

      printf("\n");
      for(int t=t_reform+1; t<=t_data_max; t++)
	{
	  printf("%d %+0.2f\n",t+1973,caldata[t]);
	  
	  if(fabs(caldata[t])>retval)
	    {
	      retval = fabs(caldata[t]);
	    }
	}
      printf("\n");
    }

  time(&stop);
  //printf("Iteration %d complete in %0.0f seconds. Max error/RMSE = %0.6f\n",iter,difftime(stop,start),retval);
  printf("Work complete in %0.0f seconds\n\n",difftime(stop,start));

  return retval;
}

double update_probs()
{
  double retval = -HUGE_VAL;
  
  // update P(NNTR-->MFN) based on pre-1980 data
  if(alt_coeffs == 0 || alt_coeffs==1 || alt_coeffs >= 3)
    {
      double avg_pre80 = 0.0;
      for(int t=3; t<t_reform; t++)
	{
	  avg_pre80 += caldata[t];
	}
      avg_pre80 = avg_pre80/((double)(t_reform-3));

      double tmp = 0.0;
      if(avg_pre80>0.0)
	{
	  tmp = tpu_trans_mat[0][1][0] * (1.0 - log(1.0+fabs(avg_pre80))*tpu_prob_update_speed);
	}
      else
	{
	  tmp = tpu_trans_mat[0][1][0] * (1.0 + log(1.0+fabs(avg_pre80))*tpu_prob_update_speed);
	}
      
      double diff = fabs(tpu_trans_mat[0][1][0]-tmp);
      if(diff>retval)
	retval=diff;
      
      tpu_trans_mat[0][1][0] = tmp;
      tpu_trans_mat[0][0][0] = 1.0-tpu_trans_mat[0][1][0];

      if(gsl_isnan(tpu_trans_mat[0][1][0]))
	{
	  printf("NaN prob detected! avg_pre80 = %0.6f\n",avg_pre80);
	}
      
      for(int t=1; t<NT; t++)
	{
	  tpu_trans_mat[0][1][t] = tpu_trans_mat[0][1][0];
	  tpu_trans_mat[0][0][t] = tpu_trans_mat[0][0][0];
	}
    }
  else
    {
      for(int t=3; t<=t_reform; t++)
	{
	  double err = caldata[t];
	  double tmp=0.0;
	  if(err<0.0)
	    {
	      tmp = tpu_trans_mat[0][1][t-1] * (1.0 + log(1.0+fabs(err)) * tpu_prob_update_speed);
	    }
	  else
	    {
	      tmp = tpu_trans_mat[0][1][t-1] * (1.0 - log(1.0+fabs(err)) * tpu_prob_update_speed);
	    }

	  double diff = fabs(tpu_trans_mat[0][1][t-1]-tmp);
	  if(diff>retval)
	    retval=diff;

	  tpu_trans_mat[0][1][t-1] = tmp;
	  tpu_trans_mat[0][0][t-1] = 1.0-tpu_trans_mat[0][1][t-1];
      
	  if(gsl_isnan(tpu_trans_mat[0][1][t-1]))
	    {
	      printf("NaN prob detected!\n");
	    }

	  for(int t=0; t<2; t++)
	    {
	      tpu_trans_mat[0][1][t] = tpu_trans_mat[0][1][2];
	      tpu_trans_mat[0][0][t] = tpu_trans_mat[0][0][2];
	    }
	  
	  for(int t=t_reform; t<NT; t++)
	    {
	      tpu_trans_mat[0][1][t] = tpu_trans_mat[0][1][t_reform-1];
	      tpu_trans_mat[0][0][t] = tpu_trans_mat[0][0][t_reform-1];
	    }
	}
    }

  // update P_t(MFN-->NNTR) from 1980-2006 based on annual NNTR gaps from 1981-2007
  for(int t=t_reform+1; t<t_data_max; t++)
    {
      double err = caldata[t];
      double tmp=0.0;
      if(err>0.0)
	{
	  tmp = tpu_trans_mat[1][0][t-1] * (1.0 + log(1.0+fabs(err)) * tpu_prob_update_speed);
	}
      else
	{
	  tmp = tpu_trans_mat[1][0][t-1] * (1.0 - log(1.0+fabs(err)) * tpu_prob_update_speed);
	}

      if(tmp<0.995 && tmp>0.005)
	{
	  double diff = fabs(tpu_trans_mat[1][0][t-1]-tmp);
	  if(diff>retval)
	    retval=diff;
	  
	  tpu_trans_mat[1][0][t-1]=tmp;
	}
      
      tpu_trans_mat[1][1][t-1] = 1.0-tpu_trans_mat[1][0][t-1];
      
      if(gsl_isnan(tpu_trans_mat[1][0][t-1]))
	{
	  printf("NaN prob detected!\n");
	}
    }
  
  
  for(int t=0; t<t_reform; t++)
    {
      tpu_trans_mat[1][0][t] = tpu_trans_mat[1][0][t_reform];
      tpu_trans_mat[1][1][t] = 1.0-tpu_trans_mat[1][0][t];
    }
  
  for(int t=t_data_max-1; t<NT; t++)
    {
      tpu_trans_mat[1][0][t] = tpu_trans_mat[1][0][t_data_max - 2];
      tpu_trans_mat[1][1][t] = tpu_trans_mat[1][1][t_data_max - 2];
    }

  return retval;
}

double update_probs_PS()
{
  double tmp = 0.0;
  
  if(caldata_PS>0)
    tmp = tpu_prob_ps * (1.0 + log(1.0+fabs(caldata_PS)) * 10*tpu_prob_update_speed);
  else
    tmp = tpu_prob_ps * (1.0 - log(1.0+fabs(caldata_PS)) * 10*tpu_prob_update_speed);

  double diff = fabs(tpu_prob_ps-tmp);  
  tpu_prob_ps = tmp;

  return diff;
}

int calibrate_probs()
{
  printf("Calibrating TPU probabilities to match estimated NNTR gap coefficients...\n");
  
  time_t start, stop;
  time(&start);

  int iter=0;
  double maxdiff = +HUGE_VAL;
  double perr = +HUGE_VAL;
  do
    {
      iter++;
      maxdiff = calc_coeffs();

      if(maxdiff<0)
	{
	  printf("Error while computing coefficients! Exiting...\n");
	  return 1;
	}
      
      if(maxdiff<coeff_err_tol)
	break;

      if(ps)
	perr = update_probs_PS();
      else
	perr = update_probs();

      printf("Iter %d. Max prob Delta = %0.6f\n",iter,perr);
	     
    }
  while(iter < max_cal_iter && perr > prob_tol);

  time(&stop);
  linebreak2();
  printf("\nCalibration complete in %0.0f seconds\n",difftime(stop,start));

  return 0;
}


/////////////////////////////////////////////////////////////////////////////
// 6. AGS sensitivity of estimation moments to parameters
/////////////////////////////////////////////////////////////////////////////

int AGS_work(double * x)
{
  time_t start2, stop2;
  time(&start2);

  if(solve_policies2_markov())
    return 1;
  
  do_trans_dyn(3);
  
  char buffer[128];

  if(perm_tpu==1)
    sprintf(buffer,"python3 -W ignore ../scripts/proc_simul.py -permtpu -ags-sens");
  else
    sprintf(buffer,"python3 -W ignore ../scripts/proc_simul.py -ags-sens");

  printf("\nProcessing simulated data: %s\n",buffer);

  if(system(buffer))
    return 1;


  if(load_caldata())
    return 1;

  int cnt=0;
  for(int t=3; t<t_data_max; t++)
    {
      x[cnt] = caldata[t];
      cnt = cnt+1;
    }

  time(&stop2);
  printf("\nProcessing complete in %0.0f seconds\n",difftime(stop2,start2)); 

  return 0;

}

int AGS_sens()
{
  linebreak();
  printf("\nComputing AGS sensitivity of estimation moments to parameter values\n");

  time_t start, stop;
  time(&start);
  
  // params:
  //// 1x Incumbent premium, maps to LR ECM elasticity
  // 1x P(NNTR --> MFN), maps to avg NTR-gap elasticity during 1970s
  // 27 x P(MFN --> NNTR) for 1980-2006, maps to NTR-gap elasticities for 1981-2007
  const int n_params = 28;
  const int n_moments = t_data_max-3; // 34

  double J[34][28] = {{0.0}};
  double x[34] = {0.0};
  double xh[34] = {0.0};

  // store values at estimated parameters
  linebreak2();
  printf("Storing baseline moments\n");
  
  if(AGS_work(x))
    return 1;

  int cnt=0;
  
  // -------------------------------------------------
  // perturb P(NNTR --> MFN)
  linebreak2();
  printf("Perturbing P(NNTR-->MFN)\n");

  double hh = 0.01/9; //hh = h + h2*tpu_trans_mat[0][1][0];

  for(int t=0; t<NT; t++)
    {
      tpu_trans_mat[0][1][t] = tpu_trans_mat[0][1][t] + hh;
      tpu_trans_mat[0][0][t] = tpu_trans_mat[0][0][t] - hh;
    }

  if(AGS_work(xh))
    return 1;

  for(int i=0; i<n_moments; i++)
    J[i][cnt] = xh[i];

  for(int t=0; t<NT; t++)
    {
      tpu_trans_mat[0][1][t] = tpu_trans_mat[0][1][t] - 2*hh;
      tpu_trans_mat[0][0][t] = tpu_trans_mat[0][0][t] + 2*hh;
    }
  
  if(AGS_work(xh))
    return 1;

  for(int i=0; i<n_moments; i++)
    J[i][cnt] = (J[i][cnt] - xh[i])/2.0;

  printf("J[:,%d] = ",cnt);
  for(int i=0; i<n_moments; i++)
    {
      printf("%0.2g ",J[i][cnt]);
    }
  printf("\n");
  
  cnt++;

  for(int t=0; t<NT; t++)
    {
      tpu_trans_mat[0][1][t] = tpu_trans_mat[0][1][t] + hh;
      tpu_trans_mat[0][0][t] = tpu_trans_mat[0][0][t] - hh;
    }

  // -------------------------------------------------
  // perturb P_t(MFN --> NNTR)
  for(int t=t_reform; t<t_data_max-1; t++)
    {
      linebreak2();
      printf("Perturbing P(MFN-->NNTR) for %d\n",t+1971);

      //hh = h + h2*tpu_trans_mat[1][0][t];
      if(t==t_data_max-2)
	hh = (1.0-Q)*0.01;
      else
	hh=0.01;

      tpu_trans_mat[1][0][t] = tpu_trans_mat[1][0][t] + hh;
      tpu_trans_mat[1][1][t] = tpu_trans_mat[1][1][t] - hh;

      if(t==t_data_max-2)
	{
	  for(int tt=t_data_max-1; tt<NT; tt++)
	    {	
	      tpu_trans_mat[1][0][tt] = tpu_trans_mat[1][0][tt] + hh;
	      tpu_trans_mat[1][1][tt] = tpu_trans_mat[1][1][tt] - hh;
      	    }
	}

      if(AGS_work(xh))
	return 1;

      for(int i=0; i<n_moments; i++)
	J[i][cnt] = xh[i];

      tpu_trans_mat[1][0][t] = tpu_trans_mat[1][0][t] - 2*hh;
      tpu_trans_mat[1][1][t] = tpu_trans_mat[1][1][t] + 2*hh;

      if(t==t_data_max-2)
	{
	  for(int tt=t_data_max-1; tt<NT; tt++)
	    {	
	      tpu_trans_mat[1][0][tt] = tpu_trans_mat[1][0][tt] - 2*hh;
	      tpu_trans_mat[1][1][tt] = tpu_trans_mat[1][1][tt] + 2*hh;
      	    }
	}

      if(AGS_work(xh))
	return 1;

      for(int i=0; i<n_moments; i++)
	J[i][cnt] = (J[i][cnt] - xh[i])/2.0;

      printf("J[:,%d] = ",cnt);
      for(int i=0; i<n_moments; i++)
	{
	  printf("%0.2g ",J[i][cnt]);
	}
      printf("\n");

      tpu_trans_mat[1][0][t] = tpu_trans_mat[1][0][t] + hh;
      tpu_trans_mat[1][1][t] = tpu_trans_mat[1][1][t] - hh;

      if(t==t_data_max-2)
	{
	  for(int tt=t_data_max-1; tt<NT; tt++)
	    {	
	      tpu_trans_mat[1][0][tt] = tpu_trans_mat[1][0][tt] + hh;
	      tpu_trans_mat[1][1][tt] = tpu_trans_mat[1][1][tt] - hh;
      	    }
	}

      cnt++;	
    }

  // -------------------------------------------------
  // write output
  FILE * file;
  if(perm_tpu==1)
    {
      
      file = fopen(OUTPUT "AGS_jacobian_permtpu.txt","w");
    }
  else
    {
      file = fopen(OUTPUT "AGS_jacobian_baseline.txt","w");
    }
  
  if(!file)
    {
      printf("Failed to open file to store Jacobian of estimation moments!\n");
      return 1;
    }

  for(int i=0; i<n_moments; i++)
    {
      for(int j=0; j<n_params; j++)
	{
	  fprintf(file,"%0.16f ",J[i][j]);
	}
      fprintf(file,"\n");
    }

  fclose(file);
  
  time(&stop);
  printf("\nProcessing complete in %0.0f seconds\n",difftime(stop,start));  

  return 0;
}

/////////////////////////////////////////////////////////////////////////////
// 7. Main function and setup/cleanup
/////////////////////////////////////////////////////////////////////////////

// process command-line arguments
int proc_cmd_args(int argc, char * argv[])
{
  for(int i=1; i<argc; i++)
    {
      if(strcmp(argv[i],"-one-sector")==0)
	{
	  multi_sector=0;
	  one_sector=1;
	  printf("Using homogeneous exporter parameters for one big sector (not avg sector)\n");
	}
      else if(strcmp(argv[i],"-avg-sector")==0)
	{
	  multi_sector=0;
	  one_sector=2;
	  printf("Using homogeneous exporter parameters for avg sector\n");
	}
      else if(strcmp(argv[i],"-calvo")==0)
	{
	  printf("Calvo model w/o endogenous exit\n");
	  calvo=1;
	  load_exporter_params_flag=0;
	}
      else if(strcmp(argv[i],"-cal-exporter-params")==0)
	{
	  printf("Calibrating exporter parameters\n");
	  cal_exporter_params_flag=1;
	}
      else if(strcmp(argv[i],"-cal-probs")==0)
	{
	  printf("Calibrating NNTR<-->MFN transition probabilities\n");
	  cal_probs_flag=1;
	}
      else if(strcmp(argv[i],"-sens-det")==0)
	{
	  printf("Doing additional deterministic versions\n");
	  sens_det=1;
	}
      else if(strcmp(argv[i],"-AGS-sens")==0)
	{
	  printf("Calculating AGS sensitivity of estimation moments\n");
	  ags=1;
	}
      else if(strcmp(argv[i],"-perm-tpu")==0)
	{
	  printf("Firms believe current TPU matrix will last forever\n");
	  perm_tpu=1;
	}
      else if(strcmp(argv[i],"-ci-lower")==0)
	{
	  printf("Using lower bound of gap-elasticity CI\n");
	  alt_coeffs=5;
	}
      else if(strcmp(argv[i],"-ci-upper")==0)
	{
	  printf("Using upper bound of gap-elasticity CI\n");
	  alt_coeffs=6;
	}
      else if(strcmp(argv[i],"-lo-xi")==0)
	{
	  printf("Smaller differences between low and high capacity exporters\n");
	  sens_xi=1;
	}
      else if(strcmp(argv[i],"-hi-xi")==0)
	{
	  printf("Larger differences between low and high capacity exporters\n");
	  sens_xi=2;
	}
      else if(strcmp(argv[i],"-lo-rhoxi")==0)
	{
	  printf("Lower iceberg trade cost persistence\n");
	  sens_xi=3;
	}
      else if(strcmp(argv[i],"-hi-rhoxi")==0)
	{
	  printf("Higher iceberg trade cost persistence\n");
	  sens_xi=4;
	}
      else if(strcmp(argv[i],"-PS-1980-temp")==0)
	{
	  printf("Fitting single reversal probability to Pierce-Schott coefficient\n");
	  printf("\tTPU from 1980 onward, WTO accession anticipated\n");
	  ps=1;
	  perm_tpu=0;
	}
      else if(strcmp(argv[i],"-PS-1980-perm")==0)
	{
	  printf("Fitting single reversal probability to Pierce-Schott coefficient\n");
	  printf("\tTPU from 1980 onward, WTO accession unanticipated\n");
	  ps=1;
	  perm_tpu=1;
	}
      else if(strcmp(argv[i],"-PS-1990-temp")==0)
	{
	  printf("Fitting single reversal probability to Pierce-Schott coefficient\n");
	  printf("\tTPU from 1990 onward, WTO accession anticipated\n");
	  ps=2;
	  perm_tpu=0;
	}
      else if(strcmp(argv[i],"-PS-1990-perm")==0)
	{
	  printf("Fitting single reversal probability to Pierce-Schott coefficient\n");
	  printf("\tTPU from 1990 onward, WTO accession unanticipated\n");
	  ps=2;
	  perm_tpu=1;
	}
      else
	{
	  return 1;
	}

      if( (ps && one_sector==1) || (ps && one_sector==2) || (ps && alt_coeffs) ||
	  (perm_tpu && one_sector==1) || (perm_tpu && one_sector==2) || (perm_tpu && alt_coeffs) )
	{
	  printf("Invalid command line arguments! Only one sensitivity analysis allowed at a time!\n");
	  return 1;
	}
    }

  return 0;
}

int setup(int argc, char * argv[])
{
  printf("Setting up model environment...\n\n");
    
  time_t start, stop;
  time(&start);

  // setup parallelization
#ifdef _OPENMP
  printf("Parallel processing with %d threads\n",omp_get_max_threads());
#else
  printf("No parallelization\n");
#endif

  perm_tpu=0;
  sunk_cost=0;
  calvo=0;
  iceberg = 0;

  if(proc_cmd_args(argc, argv))
    {
      printf("Incorrect command line argument!\n");
      return 1;
    }

  // load external data
  if(load_input_data())
    {
      printf("Failed to load calibration data!\n");
      return 1;
    }

  // load calibrated exporter params
  if(load_exporter_params_flag && load_exporter_params())
    {
      printf("Failed to load exporter params!\n");
      return 1;
    }
  
  // initialize parameters
  if(init_params())
    {
      printf("Failed to initialize parameters!\n");
      return 1;
    }

  // discretize productivity distribution
  for(int ii=0; ii<NI; ii++)
    {
      discretize_z(ii);
      calc_death_probs(ii);
    }

  for(int ii=0; ii<NI; ii++)
    {
      if(sens_xi==1)
	{
	  xi[ii] = xi[ii]*0.8;
	}
      else if(sens_xi==2)
	{
	  xi[ii] = xi[ii]*1.2;
	}
      else if(sens_xi==3)
	{
	  rho0[ii] = rho0[ii]*0.9;
	  rho1[ii] = rho1[ii]*0.9;
	}
      else if(sens_xi==4)
	{
	  rho0[ii] = rho0[ii]*1.05;
	  rho1[ii] = rho1[ii]*1.05;
	}
    }

  // initialize value functuins
  init_dp_objs_det();
  init_dp_objs_markov();
  init_dp_objs_PS();

  time(&stop);
  printf("Setup complete! Runtime = %0.0f seconds.\n",difftime(stop,start));
	  
  return 0;
}

int det_analysis()
{
  printf("Solving and simulating deterministic model...\n");

  time_t start, stop;
  time(&start);

  if(solve_policies2_det())
    return 1;
  
  do_trans_dyn(1);

  if(sens_det)
    {
      do_trans_dyn(0);
      do_trans_dyn(2);
      do_trans_dyn(4);
      do_trans_dyn(5);
    }

  time(&stop);
  printf("\nAnalysis complete! Runtime = %0.0f seconds.\n",
	 difftime(stop,start));
	  
  return 0;
}

int main(int argc, char * argv[])
{
  time_t start, stop;
  time(&start);

  // setup environment
  linebreak();    
  if(setup(argc, argv))
      return 1;

  linebreak();
  if(det_analysis())
    return 1;

  // calibration routines
  linebreak();

  if(cal_exporter_params_flag)
    {
      linebreak();
      if(calibrate_exporter_params())
	return 1;
    }
  else if(cal_probs_flag)
    {
      if(calibrate_probs())
	return 1;
    }
  // main results routine
  else
    {
      time_t start, stop;
      time(&start);

      printf("Solving model under baseline probabilities..\n");
	  
      if(ps==0)
	{
	  if(solve_policies2_markov())
	    return 1;
	}
      else
	{
	  if(solve_policies2_PS())
	    return 1;
	}
	  
      
      do_trans_dyn(3);
      time(&stop);
      printf("\nWork complete in %0.0f seconds.\n",difftime(stop,start));

    }

  if(ags)
    {
      if(AGS_sens())
	return 1;
    }
  
  // finish program
  linebreak();  
  time(&stop);
  printf("\nProgram complete! Total runtime: %0.16f seconds.\n",difftime(stop,start));

  return 0;
}
